#import <Foundation/Foundation.h>
#import <Metal/Metal.h>

#include <stdint.h>
#include <inttypes.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <math.h>
#include <float.h>
#include <fcntl.h>
#include <limits.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/sysctl.h>
#include <mach/mach.h>

#include "ds4.h"
#include "ds4_gpu.h"

/*
 * Objective-C Metal glue for the Rust engine.
 *
 * Rust owns model semantics and graph scheduling.  This file owns only
 * Metal objects: device/queue/library setup, mmap-backed weight views, command
 * batching, persistent tensors, scratch buffers, and thin wrappers around the
 * kernel files in the metal directory.  Keeping this boundary narrow makes the
 * inference path readable from Rust while still using Objective-C where Metal
 * requires it.
 */

enum {
    DS4_METAL_TENSOR_Q4_0    = 2,
    DS4_METAL_TENSOR_Q8_0    = 8,
    DS4_METAL_TENSOR_Q2_K    = 10,
    DS4_METAL_TENSOR_Q4_K    = 12,
    DS4_METAL_TENSOR_Q5_K    = 13,
    DS4_METAL_TENSOR_Q6_K    = 14,
    DS4_METAL_TENSOR_Q8_K    = 15,
    DS4_METAL_TENSOR_IQ2_XXS = 16,
};

@class DS4MetalQ4ExpertTable;

static id<MTLDevice> g_device;
static id<MTLCommandQueue> g_queue;
static id<MTLLibrary> g_library;
static id<MTLCommandBuffer> g_batch_cb;
static id<MTLComputeCommandEncoder> g_batch_enc;
static BOOL g_batch_has_work;
static NSMutableArray<id<MTLCommandBuffer>> *g_pending_cbs;
static id<MTLSharedEvent> g_selected_readback_event;
static uint64_t g_selected_readback_event_value;
static id<MTLComputePipelineState> g_set_rows_f32_i32_pipeline;
static id<MTLComputePipelineState> g_get_rows_f32_pipeline;
static id<MTLComputePipelineState> g_get_rows_f16_pipeline;
static id<MTLComputePipelineState> g_get_rows_i32_pipeline;
static id<MTLComputePipelineState> g_get_rows_q8_0_pipeline;
static id<MTLComputePipelineState> g_get_rows_q4_0_pipeline;
static id<MTLComputePipelineState> g_get_rows_q4_K_pipeline;
static id<MTLComputePipelineState> g_repeat_f32_pipeline;
static id<MTLComputePipelineState> g_concat_pipeline;
static id<MTLComputePipelineState> g_cpy_f32_f32_pipeline;
static id<MTLComputePipelineState> g_cpy_f32_f16_pipeline;
static id<MTLComputePipelineState> g_cpy_contig_f32_f16_pipeline;
static id<MTLComputePipelineState> g_cpy_f16_f32_pipeline;
static id<MTLComputePipelineState> g_cpy_f16_f16_pipeline;
static id<MTLComputePipelineState> g_cpy_contig_f16_f32_pipeline;
static id<MTLComputePipelineState> g_cpy_contig_f16_f16_pipeline;
static id<MTLComputePipelineState> g_flash_kv_stage_f16_pipeline;
static id<MTLComputePipelineState> g_swiglu_pipeline;
static id<MTLComputePipelineState> g_swiglu_flat_pipeline;
static id<MTLComputePipelineState> g_add_pipeline;
static id<MTLComputePipelineState> g_add2_pipeline;
static id<MTLComputePipelineState> g_add3_pipeline;
static id<MTLComputePipelineState> g_moe_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_sum8_pipeline;
static id<MTLComputePipelineState> g_mul_pipeline;
static id<MTLComputePipelineState> g_rms_norm_pipeline;
static id<MTLComputePipelineState> g_rms_norm_plain_pipeline;
static id<MTLComputePipelineState> g_add_rms_norm_pipeline;
static id<MTLComputePipelineState> g_rms_norm_scale_pipeline;
static id<MTLComputePipelineState> g_dsv4_qkv_rms_norm_pipeline;
static id<MTLComputePipelineState> g_hc_split_sinkhorn_pipeline;
static id<MTLComputePipelineState> g_hc_split_weighted_sum_pipeline;
static id<MTLComputePipelineState> g_hc_split_weighted_sum_norm_pipeline;
static id<MTLComputePipelineState> g_hc_weighted_sum_pipeline;
static id<MTLComputePipelineState> g_hc_weighted_sum_norm_pipeline;
static id<MTLComputePipelineState> g_output_hc_weights4_pipeline;
static id<MTLComputePipelineState> g_hc_expand_pipeline;
static id<MTLComputePipelineState> g_unary_sigmoid_pipeline;
static id<MTLComputePipelineState> g_unary_silu_pipeline;
static id<MTLComputePipelineState> g_unary_softplus_pipeline;
static id<MTLComputePipelineState> g_unary_sqrt_pipeline;
static id<MTLComputePipelineState> g_unary_clamp_pipeline;
static id<MTLComputePipelineState> g_unary_scale_pipeline;
static id<MTLComputePipelineState> g_unary_fill_pipeline;
static id<MTLComputePipelineState> g_unary_fill_f16_pipeline;
static id<MTLComputePipelineState> g_bin_mul_scalar_pipeline;
static id<MTLComputePipelineState> g_bin_div_row_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pair_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q2_k_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q2_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_pair_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group6_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group8_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group24_q4_k_id_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group24_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_q2_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_iq2_xxs_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q2_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline;
static id<MTLComputePipelineState> g_moe_stream_expert_cache_validate_pipeline;
static id<MTLComputePipelineState> g_moe_q4_gather_slots6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_table_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q4_k_sum6_pipeline;
static id<MTLArgumentEncoder> g_moe_table_q4_pair_gate_encoder;
static id<MTLArgumentEncoder> g_moe_table_q4_pair_up_encoder;
static id<MTLArgumentEncoder> g_moe_table_q4_sum_down_encoder;
static id<MTLComputePipelineState> g_rope_tail_batch_pipeline;
static id<MTLComputePipelineState> g_rope_tail_inplace_pair_pipeline;
static id<MTLComputePipelineState> g_rope_tail_inplace_pair_shared4_pipeline;
static id<MTLComputePipelineState> g_rope_tail_inplace_pair_affine_pipeline;
static id<MTLComputePipelineState> g_dsv4_fp8_kv_quantize_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_qat_pipeline;
static id<MTLComputePipelineState> g_dsv4_kv_fp8_store_pipeline;
static id<MTLComputePipelineState> g_dsv4_ratio4_shift_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_pack_ratio4_pipeline;
static id<MTLComputePipelineState> g_dsv4_softmax_pool_ratio4_direct_pipeline;
static id<MTLComputePipelineState> g_dsv4_softmax_pool_pipeline;
static id<MTLComputePipelineState> g_soft_max_f32_pipeline;
static id<MTLComputePipelineState> g_soft_max_f32_4_pipeline;
static id<MTLComputePipelineState> g_argsort_f32_i32_desc_pipeline;
static id<MTLComputePipelineState> g_argsort_merge_f32_i32_desc_pipeline;
static id<MTLComputePipelineState> g_sum_rows_f32_f32_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_scatter_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_weighted_sum_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_score_one_direct_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_store_one_pipeline;
static id<MTLComputePipelineState> g_dsv4_sort_i32_rows_asc_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_rb16_pipeline;
static id<MTLComputePipelineState> g_dsv4_softplus_sqrt_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_finalize_one_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_finalize_one_simd_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_finalize_weights_one_simd_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_weights_one_pipeline;
static id<MTLComputePipelineState> g_glm_router_select_one_pipeline;
static id<MTLComputePipelineState> g_glm_kv_lora_rms_norm_pipeline;
static id<MTLComputePipelineState> g_glm_k_b_project_pipeline;
static id<MTLComputePipelineState> g_glm_store_compact_kv_pipeline;
static id<MTLComputePipelineState> g_glm_qkv_norm_store_compact_kv_pipeline;
static id<MTLComputePipelineState> g_glm_store_indexer_k_pipeline;
static id<MTLComputePipelineState> g_glm_build_kv_cache_pipeline;
static id<MTLComputePipelineState> g_glm_build_kv_cache_decode_group4_pipeline;
static id<MTLComputePipelineState> g_glm_build_kv_cache_flash_pipeline;
static id<MTLComputePipelineState> g_glm_attention_full_pipeline;
static id<MTLComputePipelineState> g_glm_fill_selected_range_pipeline;
static id<MTLComputePipelineState> g_glm_fill_selected_range_batch_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_rope_tail_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_score_one_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_score_one_direct_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_scores_batch_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_scores_tiled_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_scores_tiled_f32_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_glm52_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_glm52_sg_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_batch_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_batch_glm52_t4_pipeline;
static id<MTLComputePipelineState> g_glm_value_project_q8_0_pipeline;
static id<MTLComputePipelineState> g_glm_value_project_q8_0_batch_heads_pipeline;
static id<MTLComputePipelineState> g_glm_value_project_q8_0_batch_heads_mma_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_partial_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_reduce_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_reduce16_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_group2_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_q2_group4_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_group8_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu2_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu4_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_addr_pair_swiglu2_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_addr_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_addr_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_addr_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_pair_swiglu_mapped_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q6_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_weights_batch_pipeline;
static id<MTLComputePipelineState> g_dsv4_hc_expand4_pipeline;
static NSMutableDictionary<NSString *, id<MTLComputePipelineState>> *g_pipeline_cache;
static NSMutableDictionary<NSString *, id<MTLBuffer>> *g_model_buffer_cache;
static NSMutableDictionary<NSString *, DS4MetalQ4ExpertTable *> *g_q4_expert_table_cache;
static NSMutableDictionary<NSString *, id> *g_q4_expert_layer_residency_cache;
static NSMutableArray<id<MTLBuffer>> *g_transient_buffers;
static id g_model_residency_set;

typedef struct {
    id<MTLBuffer> __strong mask;
    id<MTLBuffer> __strong blk;
    NSUInteger mask_bytes;
    NSUInteger blk_bytes;
    uint32_t kind;
    uint32_t n_tokens;
    uint32_t n_comp;
    uint32_t n_keys;
    uint32_t window;
    uint32_t ratio;
    uint32_t nqptg;
    uint32_t ncpsg;
    bool has_kvpad;
    bool bc_mask;
    bool valid;
    bool blk_ready;
} ds4_gpu_zero_prefix_prefill_mask_cache_entry;

enum {
    DS4_GPU_PREFILL_MASK_CACHE_RAW = 1,
    DS4_GPU_PREFILL_MASK_CACHE_RATIO4 = 2,
    DS4_GPU_PREFILL_MASK_CACHE_RATIO128 = 3,
    DS4_GPU_PREFILL_MASK_CACHE_SLOTS = 3,
};

static ds4_gpu_zero_prefix_prefill_mask_cache_entry
    g_zero_prefix_prefill_mask_cache[DS4_GPU_PREFILL_MASK_CACHE_SLOTS];
static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void);
static id<MTLBuffer> g_flash_attn_mask_buffer;
static id<MTLBuffer> g_flash_attn_zero_mask_buffer;
static id<MTLBuffer> g_flash_attn_pad_buffer;
static id<MTLBuffer> g_flash_attn_tmp_buffer;
static id<MTLBuffer> g_flash_attn_blk_buffer;
static id<MTLBuffer> g_flash_attn_ring_buffer;
static id<MTLBuffer> g_flash_attn_kv_buffer;
static id<MTLBuffer> g_glm_flash_attn_mask_buffer;
static id<MTLBuffer> g_compressor_pool_kv_buffer;
static id<MTLBuffer> g_compressor_pool_score_buffer;
static id<MTLBuffer> g_compressor_pool_score_cont_buffer;
static id<MTLBuffer> g_compressor_pool_softmax_buffer;
static id<MTLBuffer> g_compressor_pool_product_buffer;
static id<MTLBuffer> g_compressor_store_ape_buffer;
static id<MTLBuffer> g_compressor_store_score_buffer;
static id<MTLBuffer> g_embed_rows_buffer;
static id<MTLBuffer> g_router_selection_buffer;
static id<MTLBuffer> g_router_weight_sum_buffer;
static id<MTLBuffer> g_indexer_head_scores_buffer;
static id<MTLBuffer> g_indexer_topk_buffer;
static id<MTLBuffer> g_indexed_topk_buffer;
static id<MTLBuffer> g_f16_round_scratch_buffer;
static id<MTLBuffer> g_raw_store_round_buffer;
static id<MTLBuffer> g_moe_gate_scratch_buffer;
static id<MTLBuffer> g_moe_down_scratch_buffer;
static id<MTLBuffer> g_moe_id_map_buffer;
static id<MTLBuffer> g_moe_q4_gate_slots_buffer;
static id<MTLBuffer> g_moe_q4_up_slots_buffer;
static id<MTLBuffer> g_moe_q4_down_slots_buffer;
static id<MTLBuffer> g_attn_out_group_ids_buffer;
static int g_model_fd = -1;
static const void *g_model_map_ptr;
static uint64_t g_model_map_size;
static uint64_t g_model_mapped_offset;
static uint64_t g_model_mapped_size;
static uint64_t g_model_mapped_max_tensor_bytes;
static uint64_t g_tensor_alloc_live_bytes;
static uint64_t g_tensor_alloc_peak_bytes;
static pthread_mutex_t g_tensor_mu = PTHREAD_MUTEX_INITIALIZER;
static uintptr_t *g_tensor_live_slots;
static size_t g_tensor_live_cap;
static size_t g_tensor_live_count;
static size_t g_tensor_live_tombs;
static uint64_t g_model_wrap_count;
static uint64_t g_model_wrap_bytes;
static uint64_t g_model_wrap_max_bytes;
static uint64_t g_model_buffer_cache_bytes;
static uint64_t g_model_buffer_cache_evictions;
static int g_model_buffer_cache_over_limit;
static uint64_t g_stream_expert_cache_bytes;
static uint64_t g_stream_expert_cache_expert_bytes;
static uint32_t g_stream_expert_cache_entry_count;
static uint32_t g_stream_expert_cache_budget_override;
static uint64_t g_stream_expert_cache_hits;
static uint64_t g_stream_expert_cache_misses;
static uint64_t g_stream_expert_cache_evictions;
static uint64_t g_stream_expert_cache_wraps;
static uint64_t g_stream_expert_cache_clock;
static uint64_t g_stream_expert_cache_evict_advise_bytes;
static uint64_t g_stream_expert_cache_willneed_advise_bytes;
static uint64_t g_stream_expert_cache_pread_bytes;
static double g_stream_expert_cache_pread_ms;
static uint64_t g_stream_expert_cache_buffer_allocs;
static uint64_t g_stream_expert_cache_buffer_reuses;
static uint64_t g_stream_expert_cache_decode_tokens;
static uint64_t g_stream_expert_cache_hotness_decay_token;
static uint64_t g_stream_expert_timing_selected_calls;
static double g_stream_expert_timing_selected_read_ms;
static double g_stream_expert_timing_selected_sync_ms;
static double g_stream_expert_timing_selected_copy_ms;
static double g_stream_expert_timing_selected_bind_ms;
static uint64_t g_stream_expert_timing_split_layers;
static uint64_t g_stream_expert_timing_split_resident_experts;
static uint64_t g_stream_expert_timing_split_missing_experts;
static double g_stream_expert_timing_split_resident_ms;
static double g_stream_expert_timing_split_missing_ms;
static double g_stream_expert_timing_split_missing_load_ms;
static double g_stream_expert_timing_split_missing_slot_ms;
static double g_stream_expert_timing_split_missing_prune_ms;
static double g_stream_expert_timing_split_missing_addr_ms;
static double g_stream_expert_timing_split_missing_wait_ms;
static uint64_t g_stream_expert_timing_load_calls;
static double g_stream_expert_timing_load_prepare_ms;
static double g_stream_expert_timing_load_pread_ms;
static double g_stream_expert_timing_load_modify_ms;
static double g_stream_expert_timing_load_install_ms;
static uint64_t g_stream_expert_timing_prepare_batch_reuse_calls;
static double g_stream_expert_timing_prepare_batch_reuse_ms;
static uint64_t g_stream_expert_timing_prepare_buffer_calls;
static double g_stream_expert_timing_prepare_buffer_ms;
static uint64_t g_stream_expert_timing_prepare_task_experts;
static double g_stream_expert_timing_prepare_task_ms;
static uint64_t g_stream_expert_timing_reuse_scan_calls;
static uint64_t g_stream_expert_timing_reuse_scan_entries;
static double g_stream_expert_timing_reuse_scan_ms;
static double g_stream_expert_timing_reuse_clear_ms;
static uint64_t g_stream_expert_timing_readahead_calls;
static uint64_t g_stream_expert_timing_readahead_bytes;
static double g_stream_expert_timing_readahead_ms;
static uint64_t g_stream_expert_timing_cache_all_resident_layers;
static uint64_t g_stream_expert_timing_cache_all_missing_layers;
static uint64_t g_stream_expert_timing_cache_mixed_layers;
static uint64_t g_stream_expert_timing_cache_resident_experts;
static uint64_t g_stream_expert_timing_cache_missing_experts;
typedef struct {
    uint64_t selected_calls;
    double selected_read_ms;
    double selected_sync_ms;
    double selected_copy_ms;
    double selected_bind_ms;
    uint64_t split_layers;
    uint64_t split_resident_experts;
    uint64_t split_missing_experts;
    double split_resident_ms;
    double split_missing_ms;
    double split_missing_load_ms;
    double split_missing_slot_ms;
    double split_missing_prune_ms;
    double split_missing_addr_ms;
    double split_missing_wait_ms;
    uint64_t load_calls;
    double load_prepare_ms;
    double load_pread_ms;
    double load_modify_ms;
    double load_install_ms;
    uint64_t prepare_batch_reuse_calls;
    double prepare_batch_reuse_ms;
    uint64_t prepare_buffer_calls;
    double prepare_buffer_ms;
    uint64_t prepare_task_experts;
    double prepare_task_ms;
    uint64_t reuse_scan_calls;
    uint64_t reuse_scan_entries;
    double reuse_scan_ms;
    double reuse_clear_ms;
    uint64_t readahead_calls;
    uint64_t readahead_bytes;
    double readahead_ms;
    uint64_t cache_all_resident_layers;
    uint64_t cache_all_missing_layers;
    uint64_t cache_mixed_layers;
    uint64_t cache_resident_experts;
    uint64_t cache_missing_experts;
} ds4_gpu_stream_expert_timing_snapshot;
static ds4_gpu_stream_expert_timing_snapshot g_stream_expert_timing_last_report;
static int g_stream_prefill_batch_selected_addr_building;
static int g_glm_stream_expert_addr_table_building;
static uint64_t g_model_residency_count;
static int g_model_residency_added_to_queue;
static int g_glm_model_mode;
static int g_ssd_streaming_mode;
static int g_glm_streaming_prefill_full_layer_runtime;
static int g_metal4_runtime_available;
static int g_metal4_family_supported;
static int g_metal4_queue_supported;
static int g_metal4_m5_neural_accelerators_hint;
static int g_metal4_tensor_api_enabled;
static int g_metal4_tensor_api_compile_supported;
static char g_metal_device_name[128];
static int ds4_gpu_model_map_log_enabled(void);
static int ds4_gpu_stream_expert_cache_note_expert_size(
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes);
static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void);
static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats);
static void ds4_gpu_stream_expert_pending_load_clear(void);
static void ds4_gpu_stream_expert_pread_pool_shutdown(void);
static int ds4_gpu_stream_expert_timing_summary_enabled(void);
static int ds4_gpu_stream_expert_cache_entry_protected(
        uint32_t       layer,
        uint32_t       expert,
        uint32_t       protect_layer,
        const int32_t *protect_ids,
        uint32_t       n_protect);

/* The async selected-load worker registers itself so cache paths that would
 * flush/wait on command buffers (a race against the encoding thread) fail
 * the load instead; the caller then retries on the main thread. */
static pthread_t g_stream_expert_service_thread;
static int g_stream_expert_service_thread_set;

void ds4_gpu_stream_expert_cache_note_service_thread(void) {
    g_stream_expert_service_thread = pthread_self();
    g_stream_expert_service_thread_set = 1;
}

static int ds4_gpu_stream_expert_cache_on_service_thread(void) {
    return g_stream_expert_service_thread_set &&
           pthread_equal(pthread_self(), g_stream_expert_service_thread);
}
static NSUInteger g_flash_attn_mask_bytes;
static NSUInteger g_flash_attn_zero_mask_bytes;
static NSUInteger g_flash_attn_pad_bytes;
static NSUInteger g_flash_attn_tmp_bytes;
static NSUInteger g_flash_attn_blk_bytes;
static NSUInteger g_flash_attn_ring_bytes;
static NSUInteger g_flash_attn_kv_bytes;
static NSUInteger g_glm_flash_attn_mask_bytes;
static uint32_t g_glm_flash_attn_mask_pos0;
static uint32_t g_glm_flash_attn_mask_tokens;
static uint32_t g_glm_flash_attn_mask_cache_len;
static int g_glm_flash_attn_mask_valid;
static NSUInteger g_compressor_pool_kv_bytes;
static NSUInteger g_compressor_pool_score_bytes;
static NSUInteger g_compressor_pool_score_cont_bytes;
static NSUInteger g_compressor_pool_softmax_bytes;
static NSUInteger g_compressor_pool_product_bytes;
static NSUInteger g_compressor_store_ape_bytes;
static NSUInteger g_compressor_store_score_bytes;
static NSUInteger g_embed_rows_bytes;
static NSUInteger g_router_selection_bytes;
static NSUInteger g_router_weight_sum_bytes;
static NSUInteger g_indexer_head_scores_bytes;
static NSUInteger g_indexer_topk_bytes;
static NSUInteger g_indexed_topk_bytes;
static NSUInteger g_f16_round_scratch_bytes;
static NSUInteger g_raw_store_round_bytes;
static NSUInteger g_moe_gate_scratch_bytes;
static NSUInteger g_moe_down_scratch_bytes;
static NSUInteger g_moe_id_map_bytes;
static NSUInteger g_moe_q4_gate_slots_bytes;
static NSUInteger g_moe_q4_up_slots_bytes;
static NSUInteger g_moe_q4_down_slots_bytes;
static NSUInteger g_attn_out_group_ids_bytes;
static int g_initialized;
static int g_quality_mode;
static int g_mpp_invalid_env_reported;
#define DS4_METAL_MAX_ROUTED_EXPERT_USED 8
static int32_t g_routed_moe_selected_override[DS4_METAL_MAX_ROUTED_EXPERT_USED];
static uint32_t g_routed_moe_selected_override_n;
static int g_moe_selected_trace_record_initialized;
static FILE *g_moe_selected_trace_record_fp;
static uint64_t g_moe_selected_trace_record_count;
static int g_moe_selected_trace_replay_initialized;
static int32_t *g_moe_selected_trace_replay_ids;
static uint64_t g_moe_selected_trace_replay_count;
static uint64_t g_moe_selected_trace_replay_pos;

static double ds4_gpu_gib(uint64_t bytes);

static uint64_t ds4_gpu_system_memory_bytes(void) {
    uint64_t bytes = 0;
    size_t len = sizeof(bytes);
    if (sysctlbyname("hw.memsize", &bytes, &len, NULL, 0) != 0) return 0;
    return len == sizeof(bytes) ? bytes : 0;
}

static void ds4_gpu_print_device_summary(void) {
    const char *name = g_device.name ? [g_device.name UTF8String] : "unknown Metal device";
    uint64_t mem = ds4_gpu_system_memory_bytes();
    if (mem) {
        double gib = (double)mem / 1024.0 / 1024.0 / 1024.0;
        fprintf(stderr, "ds4: Metal device %s, %.2f GiB RAM\n", name, gib);
    } else {
        fprintf(stderr, "ds4: Metal device %s\n", name);
    }
}

#define DS4_METAL_MAX_MODEL_VIEWS 4096
/* Compatibility fallback for callers that cannot provide a parsed GGUF tensor
 * span. The normal DS4 engine passes the exact maximum tensor byte size. */
#define DS4_METAL_FALLBACK_MAX_TENSOR_BYTES (4ull * 1024ull * 1024ull * 1024ull)

typedef struct {
    __strong id<MTLBuffer> buffer;
    const void *model_map;
    uint64_t model_size;
    uint64_t model_offset;
    uint64_t bytes;
} ds4_gpu_model_view;

static ds4_gpu_model_view g_model_views[DS4_METAL_MAX_MODEL_VIEWS];
static uint32_t g_model_view_count;

enum {
    DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER = 80,
    DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384,
    DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = DS4_METAL_MAX_ROUTED_EXPERT_USED,
    DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES =
        DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER *
        DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT,
    DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS = 256,
    DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS = 16,
    DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS = 16,
};

typedef struct {
    uint32_t layer;
    uint32_t expert;
    uint64_t hits;
} ds4_gpu_moe_selected_hotlist_entry;

static int g_moe_selected_hotlist_initialized;
static uint64_t
    g_moe_selected_hotlist_counts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
static uint64_t g_moe_selected_hotlist_records;
static uint64_t g_moe_selected_hotlist_selections;

typedef struct {
    __strong id<MTLBuffer> gate_buffer;
    __strong id<MTLBuffer> up_buffer;
    __strong id<MTLBuffer> down_buffer;
    const void *model_map;
    uint64_t model_size;
    uint64_t gate_abs_offset;
    uint64_t up_abs_offset;
    uint64_t down_abs_offset;
    uint64_t gate_expert_bytes;
    uint64_t down_expert_bytes;
    uint64_t logical_bytes;
    uint64_t last_used;
    uint64_t use_count;
    NSUInteger gate_inner;
    NSUInteger up_inner;
    NSUInteger down_inner;
    uint64_t inflight_seq;
    uint32_t slab_slot;
    uint8_t valid;
    uint8_t slab_backed;
} ds4_gpu_stream_expert_cache_entry;

typedef struct {
    __strong id<MTLBuffer> gate_buffer;
    __strong id<MTLBuffer> up_buffer;
    __strong id<MTLBuffer> down_buffer;
    NSUInteger gate_inner;
    NSUInteger up_inner;
    NSUInteger down_inner;
} ds4_gpu_stream_expert_reusable_buffers;

static ds4_gpu_stream_expert_cache_entry
    g_stream_expert_cache[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
static ds4_gpu_stream_expert_cache_entry
    g_stream_full_expert_addr_entry[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint32_t g_stream_expert_cache_layer_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static double g_stream_expert_cache_layer_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static double g_stream_expert_cache_layer_last_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint32_t
    g_stream_expert_cache_route_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
static id<MTLBuffer> g_stream_expert_cache_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_cache_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_cache_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_cache_slabs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_start_slot[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_slot_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_slots_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_count;
static uint32_t g_stream_expert_cache_slab_total_slots;
static uint32_t g_stream_expert_cache_free_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES];
static uint32_t g_stream_expert_cache_free_slot_count;
static uint64_t g_stream_expert_cache_slab_slot_bytes;
static uint64_t g_stream_expert_cache_cb_seq;
static uint64_t g_stream_expert_cache_done_seq;
static uint64_t g_stream_expert_cache_batch_seq;
static uint64_t g_stream_expert_cache_owned_seq;
static uint64_t g_stream_expert_cache_pending_max_seq;
static id<MTLBuffer> g_stream_compact_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_compact_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_compact_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_compact_selected_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_selected_id_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_validate_status_buffer;

@interface DS4MetalTensor : NSObject
@property(nonatomic, strong) id<MTLBuffer> buffer;
@property(nonatomic, assign) uint64_t offset;
@property(nonatomic, assign) uint64_t bytes;
@property(nonatomic, assign) uint8_t owner;
@end

@implementation DS4MetalTensor
@end

@interface DS4MetalQ4ExpertTable : NSObject
@property(nonatomic, strong) id<MTLBuffer> argumentBuffer;
@property(nonatomic, strong) id<MTLBuffer> addressBuffer;
@property(nonatomic, strong) NSMutableArray<id<MTLBuffer>> *expertBuffers;
@property(nonatomic, strong) id residencySet;
@property(nonatomic, assign) BOOL residencySetAddedToQueue;
@property(nonatomic, assign) uint32_t nExpert;
@property(nonatomic, assign) uint64_t expertBytes;
@end

@implementation DS4MetalQ4ExpertTable
- (void)dealloc {
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        if (_residencySet) {
            if (_residencySetAddedToQueue &&
                g_queue &&
                [g_queue respondsToSelector:@selector(removeResidencySet:)]) {
                [g_queue removeResidencySet:_residencySet];
            }
            [_residencySet endResidency];
        }
    }
#endif
}
@end

@interface DS4MetalQ4LayerResidency : NSObject
@property(nonatomic, strong) id residencySet;
@property(nonatomic, assign) BOOL addedToQueue;
@end

@implementation DS4MetalQ4LayerResidency
- (void)dealloc {
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        if (_residencySet) {
            if (_addedToQueue &&
                g_queue &&
                [g_queue respondsToSelector:@selector(removeResidencySet:)]) {
                [g_queue removeResidencySet:_residencySet];
            }
            [_residencySet endResidency];
        }
    }
#endif
}
@end

static DS4MetalTensor *ds4_gpu_tensor_obj(ds4_gpu_tensor *tensor) {
    return (__bridge DS4MetalTensor *)tensor;
}

static const DS4MetalTensor *ds4_gpu_tensor_const_obj(const ds4_gpu_tensor *tensor) {
    return (__bridge const DS4MetalTensor *)tensor;
}

/* C code owns ds4_gpu_tensor handles as retained Objective-C objects.  Freeing
 * the same opaque handle twice would make the second __bridge_transfer release
 * an already-deallocated object, which macOS reports as malloc corruption.  The
 * live table lets free validate a handle before touching Objective-C state; the
 * same mutex also serializes the diagnostic allocation counters. */
static uint64_t ds4_gpu_tensor_ptr_hash(uintptr_t ptr) {
    uint64_t x = (uint64_t)(ptr >> 4);
    x ^= x >> 33;
    x *= UINT64_C(0xff51afd7ed558ccd);
    x ^= x >> 33;
    x *= UINT64_C(0xc4ceb9fe1a85ec53);
    x ^= x >> 33;
    return x;
}

static int ds4_gpu_tensor_live_resize_locked(size_t min_cap) {
    size_t new_cap = 1024;
    while (new_cap < min_cap) new_cap <<= 1;

    uintptr_t *new_slots = calloc(new_cap, sizeof(new_slots[0]));
    if (!new_slots) return 0;

    for (size_t i = 0; i < g_tensor_live_cap; i++) {
        const uintptr_t key = g_tensor_live_slots[i];
        if (key == 0 || key == UINTPTR_MAX) continue;

        size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (new_cap - 1);
        while (new_slots[idx] != 0) idx = (idx + 1) & (new_cap - 1);
        new_slots[idx] = key;
    }

    free(g_tensor_live_slots);
    g_tensor_live_slots = new_slots;
    g_tensor_live_cap = new_cap;
    g_tensor_live_tombs = 0;
    return 1;
}

static int ds4_gpu_tensor_live_insert_locked(const void *ptr) {
    if (!ptr || (uintptr_t)ptr == UINTPTR_MAX) return 0;
    if ((g_tensor_live_count + g_tensor_live_tombs + 1) * 10 >=
        g_tensor_live_cap * 7)
    {
        const size_t min_cap = g_tensor_live_cap ? g_tensor_live_cap * 2 : 1024;
        if (!ds4_gpu_tensor_live_resize_locked(min_cap)) return 0;
    }

    const uintptr_t key = (uintptr_t)ptr;
    size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1);
    size_t tomb = (size_t)-1;
    for (;;) {
        const uintptr_t cur = g_tensor_live_slots[idx];
        if (cur == key) return 0;
        if (cur == UINTPTR_MAX) {
            if (tomb == (size_t)-1) tomb = idx;
        } else if (cur == 0) {
            if (tomb != (size_t)-1) {
                idx = tomb;
                g_tensor_live_tombs--;
            }
            g_tensor_live_slots[idx] = key;
            g_tensor_live_count++;
            return 1;
        }
        idx = (idx + 1) & (g_tensor_live_cap - 1);
    }
}

static int ds4_gpu_tensor_live_remove_locked(const void *ptr) {
    if (!ptr || g_tensor_live_cap == 0) return 0;

    const uintptr_t key = (uintptr_t)ptr;
    size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1);
    for (;;) {
        const uintptr_t cur = g_tensor_live_slots[idx];
        if (cur == 0) return 0;
        if (cur == key) {
            g_tensor_live_slots[idx] = UINTPTR_MAX;
            g_tensor_live_count--;
            g_tensor_live_tombs++;
            return 1;
        }
        idx = (idx + 1) & (g_tensor_live_cap - 1);
    }
}

static int ds4_gpu_tensor_track_alloc_locked(
        const void *ptr,
        uint64_t bytes,
        uint64_t *live_snap,
        uint64_t *peak_snap)
{
    if (!ds4_gpu_tensor_live_insert_locked(ptr)) return 0;

    g_tensor_alloc_live_bytes += bytes;
    if (g_tensor_alloc_live_bytes > g_tensor_alloc_peak_bytes) {
        g_tensor_alloc_peak_bytes = g_tensor_alloc_live_bytes;
    }
    if (live_snap) *live_snap = g_tensor_alloc_live_bytes;
    if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes;
    return 1;
}

static int ds4_gpu_tensor_track_view_locked(const void *ptr) {
    return ds4_gpu_tensor_live_insert_locked(ptr);
}

static int ds4_gpu_tensor_prepare_free(
        ds4_gpu_tensor *tensor,
        uint8_t *owner,
        uint64_t *bytes,
        uint64_t *live_snap,
        uint64_t *peak_snap)
{
    pthread_mutex_lock(&g_tensor_mu);
    if (!ds4_gpu_tensor_live_remove_locked(tensor)) {
        pthread_mutex_unlock(&g_tensor_mu);
        fprintf(stderr,
                "ds4: Metal tensor free ignored for unknown handle %p\n",
                (void *)tensor);
        return 0;
    }

    DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor);
    const uint8_t obj_owner = obj.owner;
    const uint64_t obj_bytes = obj.bytes;
    if (obj_owner) {
        if (obj_bytes <= g_tensor_alloc_live_bytes) {
            g_tensor_alloc_live_bytes -= obj_bytes;
        } else {
            g_tensor_alloc_live_bytes = 0;
        }
    }
    if (owner) *owner = obj_owner;
    if (bytes) *bytes = obj_bytes;
    if (live_snap) *live_snap = g_tensor_alloc_live_bytes;
    if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes;
    pthread_mutex_unlock(&g_tensor_mu);
    return 1;
}

static void ds4_gpu_tensor_tracking_reset(void) {
    pthread_mutex_lock(&g_tensor_mu);
    if (g_tensor_live_count != 0) {
        fprintf(stderr,
                "ds4: Metal cleanup discarded %zu live tensor handles\n",
                g_tensor_live_count);
    }
    free(g_tensor_live_slots);
    g_tensor_live_slots = NULL;
    g_tensor_live_cap = 0;
    g_tensor_live_count = 0;
    g_tensor_live_tombs = 0;
    g_tensor_alloc_live_bytes = 0;
    g_tensor_alloc_peak_bytes = 0;
    pthread_mutex_unlock(&g_tensor_mu);
}

static id<MTLBuffer> ds4_gpu_tensor_buffer(const ds4_gpu_tensor *tensor) {
    if (!tensor) return nil;
    const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
    return obj.buffer;
}

static NSUInteger ds4_gpu_tensor_offset(const ds4_gpu_tensor *tensor) {
    if (!tensor) return 0;
    const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
    return (NSUInteger)obj.offset;
}

static id<MTLCommandBuffer> ds4_gpu_new_command_buffer(void);
static void ds4_gpu_stream_expert_cache_note_owned_created(void);

static id<MTLCommandBuffer> ds4_gpu_command_buffer(int *owned) {
    if (g_batch_cb) {
        *owned = 0;
        return g_batch_cb;
    }
    *owned = 1;
    id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
    if (cb) ds4_gpu_stream_expert_cache_note_owned_created();
    return cb;
}

static id<MTLComputeCommandEncoder> ds4_gpu_compute_encoder(id<MTLCommandBuffer> cb) {
    if (g_batch_cb && cb == g_batch_cb) {
        g_batch_has_work = YES;
        if (!g_batch_enc) g_batch_enc = [cb computeCommandEncoder];
        return g_batch_enc;
    }
    return [cb computeCommandEncoder];
}

static void ds4_gpu_end_compute_encoder(id<MTLCommandBuffer> cb, id<MTLComputeCommandEncoder> enc) {
    if (!enc) return;
    if (g_batch_cb && cb == g_batch_cb && enc == g_batch_enc) return;
    [enc endEncoding];
}

static void ds4_gpu_close_batch_encoder(void) {
    if (!g_batch_enc) return;
    [g_batch_enc endEncoding];
    g_batch_enc = nil;
}

static double g_gpu_busy_accum;
static uint64_t g_gpu_busy_cbs;

static int ds4_gpu_wait_command_buffer(id<MTLCommandBuffer> cb, const char *label) {
    [cb waitUntilCompleted];
    if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) {
        const double busy = cb.GPUEndTime - cb.GPUStartTime;
        if (busy > 0) g_gpu_busy_accum += busy;
        if ((++g_gpu_busy_cbs % 64u) == 0u) {
            fprintf(stderr, "ds4: gpu busy accum %.1f ms over %llu cbs\n",
                    g_gpu_busy_accum * 1000.0,
                    (unsigned long long)g_gpu_busy_cbs);
        }
    }
    if (cb.status == MTLCommandBufferStatusError) {
        fprintf(stderr, "ds4: Metal %s failed: %s\n",
                label, [[cb.error localizedDescription] UTF8String]);
        return 0;
    }
    return 1;
}

static id<MTLCommandBuffer> ds4_gpu_new_command_buffer(void) {
    static int initialized;
    static int use_unretained;
    if (!initialized) {
        use_unretained = getenv("DS4_METAL_UNRETAINED_COMMAND_BUFFERS") != NULL;
        initialized = 1;
    }
    if (use_unretained) {
        return [g_queue commandBufferWithUnretainedReferences];
    }
    return [g_queue commandBuffer];
}

static uint64_t ds4_gpu_exact_view_cache_limit_bytes(void) {
    static int initialized;
    static uint64_t limit_bytes;
    if (initialized) return limit_bytes;

    const uint64_t mib = 1024ull * 1024ull;
    const uint64_t gib = 1024ull * mib;
    limit_bytes = 64ull * gib;

    const char *gib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_GIB");
    if (gib_env && gib_env[0]) {
        char *end = NULL;
        unsigned long long v = strtoull(gib_env, &end, 10);
        if (end != gib_env && *end == '\0') {
            limit_bytes = v > UINT64_MAX / gib ? UINT64_MAX : (uint64_t)v * gib;
        }
    }

    const char *mib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_MIB");
    if (mib_env && mib_env[0]) {
        char *end = NULL;
        unsigned long long v = strtoull(mib_env, &end, 10);
        if (end != mib_env && *end == '\0') {
            limit_bytes = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib;
        }
    }

    initialized = 1;
    return limit_bytes;
}

static void ds4_gpu_model_buffer_cache_note_insert(uint64_t bytes) {
    if (g_model_buffer_cache_bytes > UINT64_MAX - bytes) {
        g_model_buffer_cache_bytes = UINT64_MAX;
    } else {
        g_model_buffer_cache_bytes += bytes;
    }

    const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes();
    if (limit != 0 && g_model_buffer_cache_bytes > limit) {
        g_model_buffer_cache_over_limit = 1;
    }
}

static void ds4_gpu_model_buffer_cache_clear(const char *reason) {
    if (!g_model_buffer_cache) {
        g_model_buffer_cache_bytes = 0;
        g_model_buffer_cache_over_limit = 0;
        return;
    }

    const NSUInteger entries = [g_model_buffer_cache count];
    if (entries != 0) {
        if (getenv("DS4_METAL_EXACT_VIEW_CACHE_PROFILE") != NULL) {
            fprintf(stderr,
                    "ds4: Metal exact model view cache evict reason=%s entries=%lu bytes=%.2f GiB limit=%.2f GiB\n",
                    reason ? reason : "unknown",
                    (unsigned long)entries,
                    ds4_gpu_gib(g_model_buffer_cache_bytes),
                    ds4_gpu_gib(ds4_gpu_exact_view_cache_limit_bytes()));
        }
        [g_model_buffer_cache removeAllObjects];
        g_model_buffer_cache_evictions++;
    }
    g_model_buffer_cache_bytes = 0;
    g_model_buffer_cache_over_limit = 0;
}

static void ds4_gpu_model_buffer_cache_maybe_evict(const char *reason) {
    if (g_model_buffer_cache_over_limit) {
        ds4_gpu_model_buffer_cache_clear(reason);
    }
}

static uint64_t ds4_gpu_stream_expert_cache_next_cb_seq(void) {
    if (g_stream_expert_cache_cb_seq == UINT64_MAX) {
        /*
         * A real wrap would require an astronomical number of command buffers.
         * Resetting the epoch space is still safer than letting zero become a
         * valid in-flight marker.
         */
        g_stream_expert_cache_cb_seq = 0;
        g_stream_expert_cache_done_seq = 0;
        g_stream_expert_cache_batch_seq = 0;
        g_stream_expert_cache_owned_seq = 0;
        g_stream_expert_cache_pending_max_seq = 0;
        for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) {
            for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) {
                g_stream_expert_cache[layer][expert].inflight_seq = 0;
            }
        }
    }
    return ++g_stream_expert_cache_cb_seq;
}

static void ds4_gpu_stream_expert_cache_note_batch_created(void) {
    g_stream_expert_cache_batch_seq =
        ds4_gpu_stream_expert_cache_next_cb_seq();
}

static void ds4_gpu_stream_expert_cache_note_batch_committed(void) {
    if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_pending_max_seq) {
        g_stream_expert_cache_pending_max_seq = g_stream_expert_cache_batch_seq;
    }
    g_stream_expert_cache_batch_seq = 0;
}

static void ds4_gpu_stream_expert_cache_note_owned_created(void) {
    g_stream_expert_cache_owned_seq =
        ds4_gpu_stream_expert_cache_next_cb_seq();
}

static void ds4_gpu_stream_expert_cache_note_pending_completed(void) {
    if (g_stream_expert_cache_pending_max_seq > g_stream_expert_cache_done_seq) {
        g_stream_expert_cache_done_seq = g_stream_expert_cache_pending_max_seq;
    }
    g_stream_expert_cache_pending_max_seq = 0;
}

static void ds4_gpu_stream_expert_cache_note_owned_completed(void) {
    if (g_stream_expert_cache_owned_seq > g_stream_expert_cache_done_seq) {
        g_stream_expert_cache_done_seq = g_stream_expert_cache_owned_seq;
    }
    g_stream_expert_cache_owned_seq = 0;
}

static int ds4_gpu_stream_expert_cache_entry_inflight(
        const ds4_gpu_stream_expert_cache_entry *e) {
    return e && e->valid && e->inflight_seq > g_stream_expert_cache_done_seq;
}

static int ds4_gpu_stream_expert_cache_mark_inflight(
        ds4_gpu_stream_expert_cache_entry *e) {
    if (!e || !e->valid) return 0;
    const uint64_t seq = g_stream_expert_cache_batch_seq ?
                         g_stream_expert_cache_batch_seq :
                         g_stream_expert_cache_owned_seq;
    if (seq == 0) return 0;
    e->inflight_seq = seq;
    return 1;
}

static int ds4_gpu_stream_expert_cache_mark_entries_inflight(
        ds4_gpu_stream_expert_cache_entry * const *entries,
        uint32_t n_entries,
        uint32_t active_mask) {
    if (!entries || n_entries == 0) return 0;
    for (uint32_t i = 0; i < n_entries; i++) {
        if (active_mask != 0 && (active_mask & (1u << i)) == 0) continue;
        if (!ds4_gpu_stream_expert_cache_mark_inflight(entries[i])) return 0;
    }
    return 1;
}

static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label);

static int ds4_gpu_wait_pending_command_buffers(const char *label) {
    int ok = 1;
    for (id<MTLCommandBuffer> pending in g_pending_cbs) {
        if (!ds4_gpu_wait_command_buffer(pending, label)) ok = 0;
    }
    [g_pending_cbs removeAllObjects];
    ds4_gpu_stream_expert_cache_note_pending_completed();
    if (!ok) ds4_gpu_invalidate_zero_prefix_prefill_block_maps();
    return ok;
}

static int ds4_gpu_finish_command_buffer(id<MTLCommandBuffer> cb, int owned, const char *label) {
    if (!owned) return 1;

    [cb commit];
    int ok = ds4_gpu_wait_pending_command_buffers(label);
    if (!ds4_gpu_wait_command_buffer(cb, label)) {
        ok = 0;
        ds4_gpu_invalidate_zero_prefix_prefill_block_maps();
    }
    ds4_gpu_stream_expert_cache_note_owned_completed();
    [g_transient_buffers removeAllObjects];
    ds4_gpu_model_buffer_cache_maybe_evict(label);
    return ok;
}

static int ds4_gpu_device_name_contains(const char *needle);

static int ds4_gpu_use_m5_private_scratch(void) {
    static int initialized;
    static int enabled;
    if (!initialized) {
        enabled = ds4_gpu_device_name_contains("M5");
        initialized = 1;
    }
    return enabled;
}

static int ds4_gpu_scratch_needs_cpu_access(const char *label) {
    if (!label) return 0;
    return strstr(label, "mask") != NULL ||
           strcmp(label, "ds4_attention_output_group_ids") == 0;
}

static MTLResourceOptions ds4_gpu_model_resource_options(void) {
    MTLResourceOptions options = MTLResourceStorageModeShared;
    if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) {
        options |= MTLResourceHazardTrackingModeUntracked;
    }
    return options;
}

static int ds4_gpu_ensure_scratch_buffer(
        id<MTLBuffer> __strong *buffer,
        NSUInteger    *capacity,
        NSUInteger     bytes,
        const char    *label) {
    if (*buffer && *capacity >= bytes) return 1;
    if (bytes == 0) bytes = 1;
    if (bytes > NSUIntegerMax) return 0;

    MTLResourceOptions options = MTLResourceStorageModeShared;
    if (ds4_gpu_use_m5_private_scratch() &&
        !ds4_gpu_scratch_needs_cpu_access(label)) {
        /*
         * M5 scratch buffers that only flow between Metal kernels do not need
         * CPU-visible shared storage. This reduces shared-memory traffic and
         * residency pressure for the long prefill scratch pools without
         * changing the public buffer lifetime model. Keep default hazard
         * tracking because the graph reuses these buffers across dependent
         * compute encoders.
         */
        options = MTLResourceStorageModePrivate;
    }

    *buffer = [g_device newBufferWithLength:bytes options:options];
    if (!*buffer && options != MTLResourceStorageModeShared) {
        *buffer = [g_device newBufferWithLength:bytes options:MTLResourceStorageModeShared];
    }
    if (!*buffer) {
        fprintf(stderr, "ds4: failed to allocate Metal scratch buffer %s (%llu bytes)\n",
                label, (unsigned long long)bytes);
        *capacity = 0;
        return 0;
    }
    (*buffer).label = [NSString stringWithUTF8String:label];
    *capacity = bytes;
    return 1;
}

static int ds4_gpu_ensure_zero_attention_mask(NSUInteger bytes) {
    const NSUInteger capacity = 8192u * sizeof(uint16_t);
    if (bytes > capacity) return 0;
    if (g_flash_attn_zero_mask_buffer &&
        g_flash_attn_zero_mask_bytes >= capacity) {
        return 1;
    }
    if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_zero_mask_buffer,
                                         &g_flash_attn_zero_mask_bytes,
                                         capacity,
                                         "ds4_flash_attn_zero_mask")) {
        return 0;
    }
    void *contents = [g_flash_attn_zero_mask_buffer contents];
    if (!contents) return 0;
    memset(contents, 0, g_flash_attn_zero_mask_bytes);
    return 1;
}

static uint64_t round_up_u64(uint64_t v, uint64_t align) {
    return (v + align - 1) & ~(align - 1);
}

static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint64_t max_tensor_bytes) {
    if (max_tensor_bytes != 0) return max_tensor_bytes;
    return map_size < DS4_METAL_FALLBACK_MAX_TENSOR_BYTES ?
           map_size : DS4_METAL_FALLBACK_MAX_TENSOR_BYTES;
}

static id<MTLComputePipelineState> ds4_gpu_get_pipeline(const char *function_name);
static int ds4_gpu_warm_model_views(void);
static double ds4_gpu_gib(uint64_t bytes);

static double ds4_gpu_now_ms(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0;
}

static int ds4_gpu_moe_selected_hotlist_cmp(const void *a, const void *b) {
    const ds4_gpu_moe_selected_hotlist_entry *ea = a;
    const ds4_gpu_moe_selected_hotlist_entry *eb = b;
    if (ea->hits < eb->hits) return 1;
    if (ea->hits > eb->hits) return -1;
    if (ea->layer != eb->layer) return ea->layer < eb->layer ? -1 : 1;
    if (ea->expert != eb->expert) return ea->expert < eb->expert ? -1 : 1;
    return 0;
}

static int ds4_gpu_moe_selected_hotlist_merge_requested(void) {
    return getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_MERGE") != NULL &&
           getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_FRESH") == NULL;
}

static int ds4_gpu_moe_selected_hotlist_load_existing(const char *path) {
    if (!path || !path[0] || !ds4_gpu_moe_selected_hotlist_merge_requested()) {
        return 1;
    }

    FILE *fp = fopen(path, "rb");
    if (!fp) {
        if (errno == ENOENT) return 1;
        fprintf(stderr, "ds4: failed to open selected hotlist merge file %s\n", path);
        return 0;
    }

    char line[256];
    uint64_t lineno = 0;
    uint64_t loaded_entries = 0;
    uint64_t loaded_hits = 0;
    uint64_t header_records = UINT64_MAX;
    uint64_t header_selections = UINT64_MAX;
    while (fgets(line, sizeof(line), fp)) {
        lineno++;
        char *p = line;
        while (*p && isspace((unsigned char)*p)) p++;
        if (*p == '\0') continue;
        if (*p == '#') {
            unsigned long long value = 0;
            if (sscanf(p, "# layer_records %llu", &value) == 1) {
                header_records = (uint64_t)value;
            } else if (sscanf(p, "# selections %llu", &value) == 1) {
                header_selections = (uint64_t)value;
            }
            continue;
        }

        errno = 0;
        char *end = NULL;
        unsigned long layer = strtoul(p, &end, 10);
        if (end == p || errno != 0) goto bad_line;
        p = end;
        while (*p && isspace((unsigned char)*p)) p++;

        errno = 0;
        unsigned long expert = strtoul(p, &end, 10);
        if (end == p || errno != 0) goto bad_line;
        p = end;
        while (*p && isspace((unsigned char)*p)) p++;

        errno = 0;
        unsigned long long hits = strtoull(p, &end, 10);
        if (end == p || errno != 0) goto bad_line;
        if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
            expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
            hits != 0) {
            uint64_t *dst = &g_moe_selected_hotlist_counts[layer][expert];
            if (*dst > UINT64_MAX - (uint64_t)hits) {
                *dst = UINT64_MAX;
            } else {
                *dst += (uint64_t)hits;
            }
            loaded_entries++;
            if (loaded_hits > UINT64_MAX - (uint64_t)hits) {
                loaded_hits = UINT64_MAX;
            } else {
                loaded_hits += (uint64_t)hits;
            }
        }
        continue;

bad_line:
        fprintf(stderr,
                "ds4: invalid selected hotlist merge line %" PRIu64 " in %s\n",
                lineno,
                path);
        fclose(fp);
        return 0;
    }
    if (ferror(fp)) {
        fprintf(stderr, "ds4: failed to read selected hotlist merge file %s\n", path);
        fclose(fp);
        return 0;
    }
    fclose(fp);

    g_moe_selected_hotlist_records =
        header_records != UINT64_MAX ? header_records : loaded_hits / 6u;
    g_moe_selected_hotlist_selections =
        header_selections != UINT64_MAX ? header_selections : loaded_hits;
    fprintf(stderr,
            "ds4: merged selected-id hotlist %s "
            "(%" PRIu64 " entries, %" PRIu64 " hits)\n",
            path,
            loaded_entries,
            loaded_hits);
    return 1;
}

static void ds4_gpu_moe_selected_hotlist_close(void) {
    const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST");
    if (!g_moe_selected_hotlist_initialized || !path || !path[0]) return;

    const size_t cap =
        (size_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER *
        DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
    ds4_gpu_moe_selected_hotlist_entry *entries =
        malloc(cap * sizeof(entries[0]));
    if (!entries) {
        fprintf(stderr, "ds4: failed to allocate selected hotlist entries\n");
        return;
    }

    size_t n = 0;
    for (uint32_t layer = 0;
         layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
         layer++) {
        for (uint32_t expert = 0;
             expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
             expert++) {
            const uint64_t hits =
                g_moe_selected_hotlist_counts[layer][expert];
            if (hits == 0) continue;
            entries[n++] = (ds4_gpu_moe_selected_hotlist_entry) {
                .layer = layer,
                .expert = expert,
                .hits = hits,
            };
        }
    }
    qsort(entries, n, sizeof(entries[0]), ds4_gpu_moe_selected_hotlist_cmp);

    FILE *fp = fopen(path, "wb");
    if (!fp) {
        fprintf(stderr, "ds4: failed to open selected hotlist file %s\n", path);
        free(entries);
        return;
    }
    fprintf(fp,
            "# ds4 selected-id hotlist v1\n"
            "# layer_records %" PRIu64 "\n"
            "# selections %" PRIu64 "\n"
            "# columns: layer expert hits weight\n",
            g_moe_selected_hotlist_records,
            g_moe_selected_hotlist_selections);
    for (size_t i = 0; i < n; i++) {
        fprintf(fp,
                "%u %u %" PRIu64 " 0\n",
                entries[i].layer,
                entries[i].expert,
                entries[i].hits);
    }
    free(entries);

    if (fclose(fp) != 0) {
        fprintf(stderr, "ds4: failed to close selected hotlist file %s\n", path);
    } else {
        fprintf(stderr,
                "ds4: wrote selected-id hotlist to %s "
                "(%" PRIu64 " layer records, %" PRIu64 " selections)\n",
                path,
                g_moe_selected_hotlist_records,
                g_moe_selected_hotlist_selections);
    }
}

static int ds4_gpu_moe_selected_hotlist_record(
        uint32_t       layer,
        const int32_t *selected_ids,
        uint32_t       n_selected,
        uint32_t       n_total_expert) {
    const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST");
    if (!path || !path[0]) return 1;
    if (!selected_ids ||
        n_selected == 0 ||
        n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
        return 0;
    }
    if (!g_moe_selected_hotlist_initialized) {
        g_moe_selected_hotlist_initialized = 1;
        if (!ds4_gpu_moe_selected_hotlist_load_existing(path)) return 0;
        atexit(ds4_gpu_moe_selected_hotlist_close);
    }
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 1;

    g_moe_selected_hotlist_records++;
    for (uint32_t i = 0; i < n_selected; i++) {
        if (selected_ids[i] < 0) continue;
        const uint32_t expert = (uint32_t)selected_ids[i];
        if (expert >= n_total_expert ||
            expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
            continue;
        }
        g_moe_selected_hotlist_counts[layer][expert]++;
        g_moe_selected_hotlist_selections++;
    }
    return 1;
}

static void ds4_gpu_moe_selected_trace_record_close(void) {
    if (g_moe_selected_trace_record_fp) {
        const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS");
        fclose(g_moe_selected_trace_record_fp);
        g_moe_selected_trace_record_fp = NULL;
        fprintf(stderr,
                "ds4: recorded %" PRIu64 " routed-MoE selected-id entries to %s\n",
                g_moe_selected_trace_record_count,
                path && path[0] ? path : "(unknown)");
    }
}

static int ds4_gpu_moe_selected_trace_record(
        const int32_t selected_ids[6],
        uint32_t      n_selected) {
    const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS");
    if (!path || !path[0]) return 1;
    if (n_selected != 6) {
        fprintf(stderr, "ds4: selected-id recording expects exactly 6 selected experts\n");
        return 0;
    }

    if (!g_moe_selected_trace_record_initialized) {
        g_moe_selected_trace_record_initialized = 1;
        g_moe_selected_trace_record_fp = fopen(path, "wb");
        if (!g_moe_selected_trace_record_fp) {
            fprintf(stderr, "ds4: failed to open selected-id record file %s\n", path);
            return 0;
        }
        setvbuf(g_moe_selected_trace_record_fp, NULL, _IOFBF, 1u << 20);
        atexit(ds4_gpu_moe_selected_trace_record_close);
    }

    if (fwrite(selected_ids, sizeof(selected_ids[0]), n_selected, g_moe_selected_trace_record_fp) != n_selected) {
        fprintf(stderr, "ds4: failed to write selected-id record file %s\n", path);
        return 0;
    }
    if (fflush(g_moe_selected_trace_record_fp) != 0) {
        fprintf(stderr, "ds4: failed to flush selected-id record file %s\n", path);
        return 0;
    }
    g_moe_selected_trace_record_count++;
    return 1;
}

static int ds4_gpu_moe_selected_trace_replay(
        int32_t  selected_ids[6],
        uint32_t n_selected) {
    const char *path = getenv("DS4_MOE_REPLAY_SELECTED_IDS");
    if (!path || !path[0]) return 0;
    if (n_selected != 6) {
        fprintf(stderr, "ds4: selected-id replay expects exactly 6 selected experts\n");
        return -1;
    }

    if (!g_moe_selected_trace_replay_initialized) {
        g_moe_selected_trace_replay_initialized = 1;
        FILE *fp = fopen(path, "rb");
        if (!fp) {
            fprintf(stderr, "ds4: failed to open selected-id replay file %s\n", path);
            return -1;
        }
        if (fseeko(fp, 0, SEEK_END) != 0) {
            fprintf(stderr, "ds4: failed to seek selected-id replay file %s\n", path);
            fclose(fp);
            return -1;
        }
        const off_t end = ftello(fp);
        if (end < 0) {
            fprintf(stderr, "ds4: failed to size selected-id replay file %s\n", path);
            fclose(fp);
            return -1;
        }
        if (fseeko(fp, 0, SEEK_SET) != 0) {
            fprintf(stderr, "ds4: failed to rewind selected-id replay file %s\n", path);
            fclose(fp);
            return -1;
        }

        const uint64_t bytes = (uint64_t)end;
        const uint64_t entry_bytes = (uint64_t)n_selected * sizeof(selected_ids[0]);
        if (bytes == 0 || (bytes % entry_bytes) != 0) {
            fprintf(stderr,
                    "ds4: selected-id replay file %s has invalid size %" PRIu64 "\n",
                    path,
                    bytes);
            fclose(fp);
            return -1;
        }
        if (bytes > SIZE_MAX) {
            fprintf(stderr, "ds4: selected-id replay file %s is too large\n", path);
            fclose(fp);
            return -1;
        }
        g_moe_selected_trace_replay_count = bytes / entry_bytes;
        g_moe_selected_trace_replay_ids = malloc((size_t)bytes);
        if (!g_moe_selected_trace_replay_ids) {
            fprintf(stderr, "ds4: failed to allocate selected-id replay buffer\n");
            fclose(fp);
            return -1;
        }
        if (fread(g_moe_selected_trace_replay_ids, 1, (size_t)bytes, fp) != (size_t)bytes) {
            fprintf(stderr, "ds4: failed to read selected-id replay file %s\n", path);
            fclose(fp);
            free(g_moe_selected_trace_replay_ids);
            g_moe_selected_trace_replay_ids = NULL;
            return -1;
        }
        fclose(fp);
        fprintf(stderr,
                "ds4: loaded %" PRIu64 " routed-MoE selected-id entries from %s\n",
                g_moe_selected_trace_replay_count,
                path);
    }

    if (g_moe_selected_trace_replay_pos >= g_moe_selected_trace_replay_count) {
        fprintf(stderr,
                "ds4: selected-id replay exhausted after %" PRIu64 " entries\n",
                g_moe_selected_trace_replay_pos);
        return -1;
    }
    memcpy(selected_ids,
           g_moe_selected_trace_replay_ids + g_moe_selected_trace_replay_pos * n_selected,
           (size_t)n_selected * sizeof(selected_ids[0]));
    g_moe_selected_trace_replay_pos++;
    return 1;
}

static int ds4_gpu_progress_enabled(void) {
    return ds4_log_is_tty(stderr);
}

static void ds4_gpu_progress_begin(const char *what) {
    if (!ds4_gpu_progress_enabled()) return;
    fprintf(stderr, "ds4: %s...", what);
    fflush(stderr);
}

static void ds4_gpu_progress_done(void) {
    if (!ds4_gpu_progress_enabled()) return;
    fputs(" done\n", stderr);
    fflush(stderr);
}

static void ds4_gpu_progress_failed(void) {
    if (!ds4_gpu_progress_enabled()) return;
    fputs(" failed\n", stderr);
    fflush(stderr);
}

static void ds4_gpu_model_views_clear(void) {
    for (uint32_t i = 0; i < g_model_view_count; i++) {
        g_model_views[i].buffer = nil;
        g_model_views[i].model_map = NULL;
        g_model_views[i].model_size = 0;
        g_model_views[i].model_offset = 0;
        g_model_views[i].bytes = 0;
    }
    g_model_view_count = 0;
}

static void ds4_gpu_model_residency_clear(void) {
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        if (g_model_residency_set) {
            if (g_model_residency_added_to_queue &&
                g_queue &&
                [g_queue respondsToSelector:@selector(removeResidencySet:)]) {
                [g_queue removeResidencySet:g_model_residency_set];
            }
            [g_model_residency_set endResidency];
            [g_model_residency_set removeAllAllocations];
            g_model_residency_set = nil;
        }
    }
#endif
    g_model_residency_count = 0;
    g_model_residency_added_to_queue = 0;
}

/* TP sharding keeps only this rank's expert ranges warm,
 * so whole-view residency requests (which would page in the full file)
 * must be skipped; pages fault in lazily through the same view buffers,
 * exactly like ssd-streaming mode. */
static int g_model_residency_skipped;

void ds4_gpu_model_residency_skip(int skip) {
    g_model_residency_skipped = skip;
}

static int ds4_gpu_model_residency_request_views(void) {
    if (g_model_view_count == 0 ||
        g_ssd_streaming_mode ||
        g_model_residency_skipped ||
        getenv("DS4_METAL_NO_RESIDENCY") != NULL) {
        return 1;
    }

#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        /*
         * Register all model views as one residency set before inference. This
         * is a GPU residency/budgeting hint, not a request to fault the whole
         * 80+ GB file into memory. Its purpose is to make the driver see the
         * complete set of large shared allocations during setup instead of
         * discovering them lazily from the first measured graph command, where
         * VM validation and residency accounting would look like model compute.
         */
        MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init];
        desc.label = @"ds4_model";
        desc.initialCapacity = g_model_view_count;

        NSError *error = nil;
        g_model_residency_set = [g_device newResidencySetWithDescriptor:desc error:&error];
        if (!g_model_residency_set) {
            fprintf(stderr, "ds4: Metal model residency set creation failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            return 0;
        }

        for (uint32_t i = 0; i < g_model_view_count; i++) {
            [g_model_residency_set addAllocation:g_model_views[i].buffer];
        }
        [g_model_residency_set commit];
        [g_model_residency_set requestResidency];
        if (getenv("DS4_METAL_DISABLE_QUEUE_RESIDENCY_SET") == NULL &&
            g_queue &&
            [g_queue respondsToSelector:@selector(addResidencySet:)]) {
            [g_queue addResidencySet:g_model_residency_set];
            g_model_residency_added_to_queue = 1;
        }
        g_model_residency_count = g_model_view_count;
    }
#endif

    return 1;
}

static int ds4_gpu_add_model_view_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    map_offset,
        uint64_t    map_size,
        uint64_t    max_tensor_bytes,
        bool        use_default_view_cap,
        uint64_t   *mapped_model_size_out) {
    const uint64_t page = (uint64_t)getpagesize();
    const uintptr_t model_addr = (uintptr_t)model_map;

    if ((model_addr & (uintptr_t)(page - 1)) != 0) {
        fprintf(stderr, "ds4: Metal model mmap base is not page aligned\n");
        return 0;
    }
    if (map_offset > model_size || map_size > model_size - map_offset) {
        fprintf(stderr, "ds4: Metal model mapped range is outside the GGUF mapping\n");
        return 0;
    }
    const uint64_t page_model_offset = map_offset & ~(page - 1);
    const uint64_t leading = map_offset - page_model_offset;
    if (map_size > UINT64_MAX - leading ||
        leading + map_size > UINT64_MAX - (page - 1))
    {
        fprintf(stderr, "ds4: Metal model mapped range overflows page alignment\n");
        return 0;
    }
    const uint64_t mapped_model_size = round_up_u64(leading + map_size, page);
    uint64_t max_buffer = (uint64_t)[g_device maxBufferLength];
    max_buffer &= ~(page - 1);

    /*
     * Wrap only the tensor-data part of the GGUF file. Metadata is parsed by the
     * CPU and is never dereferenced by kernels, so exposing it to Metal only
     * grows the residency set and the VM range the driver must validate.
     *
     * Metal buffers have a device-specific maximum length, and this model is
     * larger than that maximum on the target machines. Creating one no-copy
     * buffer per tensor would avoid the length limit, but it would also move a
     * lot of VM-object creation and residency bookkeeping into graph setup. The
     * stable shape here is a tiny number of page-aligned views created once.
     *
     * Adjacent views intentionally overlap by more than the largest tensor, plus
     * one page for alignment. That invariant guarantees every tensor lies wholly
     * inside at least one view, so hot paths pass one buffer and one inner byte
     * offset. We never split a weight tensor across command encoders.
     */
    if (max_tensor_bytes > map_size) {
        fprintf(stderr, "ds4: Metal model max tensor span is larger than a mapped tensor span\n");
        return 0;
    }
    if (max_tensor_bytes > UINT64_MAX - (page - 1)) {
        fprintf(stderr, "ds4: Metal model max tensor span overflows page alignment\n");
        return 0;
    }
    const uint64_t max_tensor_rounded = round_up_u64(max_tensor_bytes, page);
    if (max_tensor_rounded > UINT64_MAX - page) {
        fprintf(stderr, "ds4: Metal model view overlap overflows page slack\n");
        return 0;
    }
    const uint64_t overlap = max_tensor_rounded + page;
    if (max_buffer == 0 || max_buffer <= overlap) {
        fprintf(stderr,
                "ds4: Metal maxBufferLength is too small for DS4 model views "
                "(max tensor %.2f GiB, max buffer %.2f GiB)\n",
                ds4_gpu_gib(max_tensor_bytes),
                ds4_gpu_gib(max_buffer));
        return 0;
    }

    uint64_t view_limit = max_buffer;
    const char *view_limit_env = getenv("DS4_METAL_MODEL_VIEW_MAX_GIB");
    if (view_limit_env && view_limit_env[0]) {
        char *end = NULL;
        unsigned long long gib = strtoull(view_limit_env, &end, 10);
        if (end != view_limit_env && gib > 0) {
            uint64_t env_limit = gib * 1024ull * 1024ull * 1024ull;
            env_limit &= ~(page - 1);
            if (env_limit > 0) view_limit = env_limit;
        }
    } else if (use_default_view_cap && mapped_model_size > max_buffer) {
        /*
         * Very large no-copy buffers can make Metal's VM validation dominate
         * startup or the first graph command on multi-hundred-GiB slices. Keep
         * ordinary contiguous model mappings unchanged, but let distributed
         * span maps use smaller overlapping views when a range already has to
         * be split.
         */
        const uint64_t default_limit = 128ull * 1024ull * 1024ull * 1024ull;
        if (view_limit > default_limit) view_limit = default_limit;
    }
    if (view_limit > max_buffer) view_limit = max_buffer;
    view_limit &= ~(page - 1);
    if (view_limit == 0 || view_limit <= overlap) {
        fprintf(stderr,
                "ds4: Metal model view cap is too small for DS4 model views "
                "(cap %.2f GiB, max tensor %.2f GiB)\n",
                ds4_gpu_gib(view_limit),
                ds4_gpu_gib(max_tensor_bytes));
        return 0;
    }

    const uint64_t step = view_limit - overlap;
    uint64_t off = 0;
    while (off < mapped_model_size) {
        if (g_model_view_count == DS4_METAL_MAX_MODEL_VIEWS) {
            fprintf(stderr, "ds4: Metal model needs more mapped views than expected\n");
            return 0;
        }

        uint64_t view_bytes = mapped_model_size - off;
        if (view_bytes > view_limit) view_bytes = view_limit;

        id<MTLBuffer> buffer = [g_device newBufferWithBytesNoCopy:(void *)(model_addr + page_model_offset + off)
                                                           length:(NSUInteger)view_bytes
                                                          options:ds4_gpu_model_resource_options()
                                                      deallocator:nil];
        if (!buffer) {
            fprintf(stderr,
                    "ds4: Metal could not wrap mmaped model view at %.2f GiB, size %.2f GiB\n",
                    (double)(page_model_offset + off) / (1024.0 * 1024.0 * 1024.0),
                    (double)view_bytes / (1024.0 * 1024.0 * 1024.0));
            return 0;
        }
        buffer.label = [NSString stringWithFormat:@"ds4_model_view_%u", g_model_view_count];

        g_model_views[g_model_view_count].buffer = buffer;
        g_model_views[g_model_view_count].model_map = model_map;
        g_model_views[g_model_view_count].model_size = model_size;
        g_model_views[g_model_view_count].model_offset = page_model_offset + off;
        g_model_views[g_model_view_count].bytes = view_bytes;
        g_model_view_count++;

        g_model_wrap_count++;
        g_model_wrap_bytes += view_bytes;
        if (view_bytes > g_model_wrap_max_bytes) g_model_wrap_max_bytes = view_bytes;

        if (off + view_bytes >= mapped_model_size) break;
        off += step;
    }

    if (mapped_model_size_out) *mapped_model_size_out += mapped_model_size;
    return 1;
}

static int ds4_gpu_finish_model_views(
        double t0,
        uint64_t mapped_model_size,
        uint64_t display_offset) {
    const double t_mapped = ds4_gpu_now_ms();
    const int request_residency =
        !g_ssd_streaming_mode &&
        getenv("DS4_METAL_NO_RESIDENCY") == NULL;
    if (request_residency) ds4_gpu_progress_begin("requesting Metal residency (may take tens of seconds)");
    if (!ds4_gpu_model_residency_request_views()) {
        if (request_residency) ds4_gpu_progress_failed();
        return 0;
    }
    if (request_residency) ds4_gpu_progress_done();
    const double t_resident = ds4_gpu_now_ms();
    int warmed = 1;
    const double t_warm0 = ds4_gpu_now_ms();
    const int warm_model_views = !g_ssd_streaming_mode &&
                                 getenv("DS4_METAL_NO_RESIDENCY") == NULL &&
                                 getenv("DS4_METAL_NO_MODEL_WARMUP") == NULL;
    if (warm_model_views) {
        /*
         * The first GPU command touching no-copy mmap storage can pay command
         * queue setup, page-table validation, and shared-allocation residency
         * costs. Sample each model view here so timed graph execution starts
         * after that one-time work. The stride is intentionally coarse: this is
         * a validation touch over the VM ranges, not a full model prefetch. A
         * dense prefetch would create exactly the kind of memory pressure and
         * startup stalls this path is designed to avoid.
         */
        if (g_model_residency_skipped) {
            /* TP sharding: a single command buffer binding every
             * view demands residency of them all and OOMs; the engine's
             * CPU-side sharded warm pre-faults the owned bytes instead. */
            warmed = 1;
        } else {
            ds4_gpu_progress_begin("warming Metal model views");
            warmed = ds4_gpu_warm_model_views();
            if (warmed) ds4_gpu_progress_done();
            else ds4_gpu_progress_failed();
        }
    }
    const double t_warm = ds4_gpu_now_ms();
    if (ds4_gpu_model_map_log_enabled()) {
        fprintf(stderr,
                "ds4: Metal model views created in %.3f ms, residency requested in %.3f ms, warmup %.3f ms (mapped %.2f MiB from offset %.2f MiB)\n",
                t_mapped - t0,
                t_resident - t_mapped,
                t_warm - t_warm0,
                mapped_model_size / 1024.0 / 1024.0,
                display_offset / 1024.0 / 1024.0);
    }
    if (!warmed) return 0;
    return 1;
}

static int ds4_gpu_map_model_views(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    map_offset,
        uint64_t    map_size,
        uint64_t    max_tensor_bytes) {
    const double t0 = ds4_gpu_now_ms();
    uint64_t mapped_model_size = 0;
    if (!ds4_gpu_add_model_view_range(model_map,
                                      model_size,
                                      map_offset,
                                      map_size,
                                      max_tensor_bytes,
                                      false,
                                      &mapped_model_size)) {
        return 0;
    }
    return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset);
}

static id<MTLBuffer> ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) {
    if (bytes == 0) bytes = 1;

    id<MTLBuffer> buffer = [g_device newBufferWithLength:bytes
                                                 options:MTLResourceStorageModeShared];
    if (!buffer) {
        fprintf(stderr, "ds4: failed to allocate Metal transient buffer %s (%llu bytes)\n",
                label ? label : "(unnamed)", (unsigned long long)bytes);
        return nil;
    }
    if (label) buffer.label = [NSString stringWithUTF8String:label];

    /*
     * CPU-filled buffers must survive until their command buffer completes.
     * A local ObjC strong variable is not enough when the encoder function
     * returns before the caller commits the command buffer.
     */
    [g_transient_buffers addObject:buffer];
    return buffer;
}

static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) {
    if (getenv("DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL ||
        getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) {
        return 0;
    }
    return ds4_gpu_device_name_contains("M3") ||
           getenv("DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL;
}

static ds4_gpu_zero_prefix_prefill_mask_cache_entry *
ds4_gpu_get_zero_prefix_prefill_mask_cache(
        uint32_t   kind,
        uint32_t   n_tokens,
        uint32_t   n_comp,
        uint32_t   n_keys,
        uint32_t   window,
        uint32_t   ratio,
        uint32_t   nqptg,
        uint32_t   ncpsg,
        bool       has_kvpad,
        bool       bc_mask,
        NSUInteger mask_bytes,
        NSUInteger blk_bytes,
        bool      *created) {
    if (created) *created = false;
    if (!created || !ds4_gpu_zero_prefix_prefill_mask_cache_enabled() ||
        mask_bytes == 0 || blk_bytes == 0) {
        return NULL;
    }

    uint32_t slot = UINT32_MAX;
    switch (kind) {
        case DS4_GPU_PREFILL_MASK_CACHE_RAW:      slot = 0; break;
        case DS4_GPU_PREFILL_MASK_CACHE_RATIO4:   slot = 1; break;
        case DS4_GPU_PREFILL_MASK_CACHE_RATIO128: slot = 2; break;
        default: return NULL;
    }

    ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry =
        &g_zero_prefix_prefill_mask_cache[slot];
    if (entry->valid && entry->mask && entry->blk &&
        entry->mask_bytes == mask_bytes && entry->blk_bytes == blk_bytes &&
        entry->kind == kind && entry->n_tokens == n_tokens &&
        entry->n_comp == n_comp && entry->n_keys == n_keys &&
        entry->window == window && entry->ratio == ratio &&
        entry->nqptg == nqptg && entry->ncpsg == ncpsg &&
        entry->has_kvpad == has_kvpad && entry->bc_mask == bc_mask) {
        return entry;
    }

    id<MTLBuffer> mask = [g_device newBufferWithLength:mask_bytes
                                                options:MTLResourceStorageModeShared];
    id<MTLBuffer> blk = [g_device newBufferWithLength:blk_bytes
                                               options:MTLResourceStorageModePrivate];
    if (!blk) {
        blk = [g_device newBufferWithLength:blk_bytes
                                    options:MTLResourceStorageModeShared];
    }
    if (!mask || !blk) return NULL;

    mask.label = [NSString stringWithFormat:@"ds4_prefill_mask_cache_%u", kind];
    blk.label = [NSString stringWithFormat:@"ds4_prefill_blk_cache_%u", kind];

    /* A replaced entry may still be referenced by an uncommitted or in-flight
     * command buffer. Keep its resources alive with the other batch-scoped
     * buffers until command completion instead of mutating or releasing them. */
    if (entry->mask) [g_transient_buffers addObject:entry->mask];
    if (entry->blk) [g_transient_buffers addObject:entry->blk];

    entry->mask = mask;
    entry->blk = blk;
    entry->mask_bytes = mask_bytes;
    entry->blk_bytes = blk_bytes;
    entry->kind = kind;
    entry->n_tokens = n_tokens;
    entry->n_comp = n_comp;
    entry->n_keys = n_keys;
    entry->window = window;
    entry->ratio = ratio;
    entry->nqptg = nqptg;
    entry->ncpsg = ncpsg;
    entry->has_kvpad = has_kvpad;
    entry->bc_mask = bc_mask;
    /* The caller publishes the entry only after synchronously filling mask. */
    entry->valid = false;
    entry->blk_ready = false;
    *created = true;
    return entry;
}

static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void) {
    for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) {
        g_zero_prefix_prefill_mask_cache[i].blk_ready = false;
    }
}

static void ds4_gpu_clear_zero_prefix_prefill_mask_cache(void) {
    for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) {
        ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry =
            &g_zero_prefix_prefill_mask_cache[i];
        entry->mask = nil;
        entry->blk = nil;
        entry->mask_bytes = 0;
        entry->blk_bytes = 0;
        entry->kind = 0;
        entry->n_tokens = 0;
        entry->n_comp = 0;
        entry->n_keys = 0;
        entry->window = 0;
        entry->ratio = 0;
        entry->nqptg = 0;
        entry->ncpsg = 0;
        entry->has_kvpad = false;
        entry->bc_mask = false;
        entry->valid = false;
        entry->blk_ready = false;
    }
}

void ds4_gpu_release_zero_prefix_prefill_mask_cache(void) {
    /* Layer-major prefill waits each layer before advancing, so its final
     * release point has no outstanding cache users. Keep the guard here for
     * diagnostic callers that may have an open or asynchronously flushed CB. */
    if (!g_initialized || g_batch_cb ||
        (g_pending_cbs && [g_pending_cbs count] != 0)) {
        return;
    }
    ds4_gpu_clear_zero_prefix_prefill_mask_cache();
}

static id<MTLComputePipelineState> ds4_gpu_get_mul_mm_pipeline(
        const char *function_name,
        bool        bc_inp,
        bool        bc_out) {
    NSString *key = [NSString stringWithFormat:@"%s_bci=%d_bco=%d",
                     function_name, bc_inp ? 1 : 0, bc_out ? 1 : 0];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700];
    [constants setConstantValue:&bc_out type:MTLDataTypeBool atIndex:701];

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_mul_mm_id_pipeline(
        const char *function_name,
        bool        bc_inp) {
    NSString *key = [NSString stringWithFormat:@"%s_bci=%d",
                     function_name, bc_inp ? 1 : 0];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700];

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_pipeline(
        const char *function_name) {
    NSString *key = [NSString stringWithFormat:@"%s", function_name];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found\n", function_name);
        return nil;
    }

    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static int ds4_gpu_disable_hot_pipeline_statics(void) {
    static int initialized;
    static int disabled;
    if (!initialized) {
        disabled = getenv("DS4_METAL_DISABLE_HOT_PIPELINE_STATICS") != NULL;
        initialized = 1;
    }
    return disabled;
}

static id<MTLComputePipelineState> ds4_gpu_hot_pipeline(
        id<MTLComputePipelineState> pipeline,
        const char *fallback_name) {
    if (!ds4_gpu_disable_hot_pipeline_statics()) return pipeline;
    return ds4_gpu_get_pipeline(fallback_name);
}

static int ds4_gpu_use_compressor_pair_nr4(void) {
    static int initialized;
    static int enabled;
    if (!initialized) {
        enabled = getenv("DS4_METAL_COMPRESSOR_PAIR_NR4") != NULL;
        initialized = 1;
    }
    return enabled;
}

static int ds4_gpu_device_name_contains(const char *needle);

static int ds4_gpu_env_value_eq(const char *v, size_t n, const char *literal) {
    size_t m = strlen(literal);
    if (n != m) return 0;
    for (size_t i = 0; i < n; i++) {
        if (tolower((unsigned char)v[i]) != tolower((unsigned char)literal[i])) return 0;
    }
    return 1;
}

static int ds4_gpu_env_bool(const char *name) {
    const char *v = getenv(name);
    if (!v) return -1;

    while (isspace((unsigned char)*v)) v++;
    size_t n = strlen(v);
    while (n > 0 && isspace((unsigned char)v[n - 1])) n--;
    if (n == 0) return 1;

    if (ds4_gpu_env_value_eq(v, n, "1") ||
        ds4_gpu_env_value_eq(v, n, "true") ||
        ds4_gpu_env_value_eq(v, n, "yes") ||
        ds4_gpu_env_value_eq(v, n, "on")) {
        return 1;
    }
    if (ds4_gpu_env_value_eq(v, n, "0") ||
        ds4_gpu_env_value_eq(v, n, "false") ||
        ds4_gpu_env_value_eq(v, n, "no") ||
        ds4_gpu_env_value_eq(v, n, "off")) {
        return 0;
    }

    if (!g_mpp_invalid_env_reported) {
        fprintf(stderr,
                "ds4: invalid Metal boolean environment value %s=%.*s; treating presence as enabled\n",
                name, (int)n, v);
        g_mpp_invalid_env_reported = 1;
    }
    return 1;
}

static uint64_t ds4_gpu_env_u64(const char *name,
                                uint64_t    fallback,
                                uint64_t    min_value,
                                uint64_t    max_value) {
    const char *v = getenv(name);
    if (!v) return fallback;
    while (isspace((unsigned char)*v)) v++;
    if (!*v) return fallback;

    errno = 0;
    char *end = NULL;
    unsigned long long parsed = strtoull(v, &end, 10);
    if (end == v || errno == ERANGE) return fallback;
    while (isspace((unsigned char)*end)) end++;
    if (*end) return fallback;

    if (parsed < min_value) return fallback;
    uint64_t value = (uint64_t)parsed;
    if (value > max_value) value = max_value;
    return value;
}

static uint32_t ds4_gpu_glm_full_attention_max_cache_len(void) {
    /*
     * kernel_glm_attention_full stores one score per visible token in
     * threadgroup memory, plus 256 reduction slots.  Keep the default under
     * the 32 KiB envelope used by current Apple GPUs while allowing short
     * dense decode to step past the 4096-token prefill boundary.
     */
    return 7680u;
}

static uint32_t ds4_gpu_glm_flash_attention_max_cache_len(void) {
    /*
     * Staged FlashAttention uses fixed-size threadgroup scratch and separate
     * KV staging buffers, so it is not bound by the legacy full-attention
     * kernel's one-score-per-token threadgroup-memory envelope.
     */
    return 8192u;
}

static int ds4_gpu_mpp_available(void) {
    return g_metal4_tensor_api_enabled && !g_quality_mode;
}

/*
 * Retained Metal4 defaults live here instead of behind user-visible options.
 * The public runtime has one automatic accelerated path plus the global
 * DS4_METAL_DISABLE_METAL4 comparison switch.  Benchmark-only alternatives that
 * lost during M5 work are removed or kept out of the dispatch path so future
 * changes do not accidentally turn old experiments into new modes.
 */
static int ds4_gpu_use_mpp_attn_out_low_matmul(void) {
    return ds4_gpu_mpp_available();
}

enum {
    DS4_METAL_ATTN_OUT_MPP_TILE_N = 64,
};

static void ds4_gpu_warn_mpp_fallback(void) {
    static int warned;
    if (!warned) {
        fprintf(stderr, "ds4: accelerated Metal prefill matmul unavailable; falling back to legacy kernel\n");
        warned = 1;
    }
}

static int ds4_gpu_device_name_contains(const char *needle) {
    return g_metal_device_name[0] != '\0' && strstr(g_metal_device_name, needle) != NULL;
}

static int ds4_gpu_compile_tensor_probe(void) {
#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000
    if (!g_device) return 0;
    if (@available(macOS 26.0, *)) {
        const char *src =
            "#include <metal_stdlib>\n"
            "#include <metal_tensor>\n"
            "#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>\n"
            "using namespace metal;\n"
            "using namespace mpp::tensor_ops;\n"
            "kernel void ds4_tensor_probe(\n"
            "        tensor<device half,  dextents<int32_t, 2>> A [[buffer(0)]],\n"
            "        tensor<device half,  dextents<int32_t, 2>> B [[buffer(1)]],\n"
            "        device float *C [[buffer(2)]],\n"
            "        uint2 tgid [[threadgroup_position_in_grid]]) {\n"
            "    auto tA = A.slice(0, (int)tgid.y);\n"
            "    auto tB = B.slice((int)tgid.x, 0);\n"
            "    matmul2d<matmul2d_descriptor(16, 16, dynamic_extent), execution_simdgroups<4>> mm;\n"
            "    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();\n"
            "    auto sA = tA.slice(0, 0);\n"
            "    auto sB = tB.slice(0, 0);\n"
            "    mm.run(sB, sA, cT);\n"
            "    auto tC = tensor<device float, dextents<int32_t, 2>, tensor_inline>(C, dextents<int32_t, 2>(16, 16));\n"
            "    cT.store(tC);\n"
            "}\n";

        NSError *error = nil;
        NSString *source = [NSString stringWithUTF8String:src];
        id<MTLLibrary> probe_library = [g_device newLibraryWithSource:source options:[MTLCompileOptions new] error:&error];
        if (!probe_library) {
            fprintf(stderr, "ds4: Metal 4 tensor API probe compile failed: %s\n",
                    error ? [[error localizedDescription] UTF8String] : "(unknown)");
            return 0;
        }
        id<MTLFunction> fn = [probe_library newFunctionWithName:@"ds4_tensor_probe"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal 4 tensor API probe function missing\n");
            return 0;
        }
        error = nil;
        id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!pipeline) {
            fprintf(stderr, "ds4: Metal 4 tensor API probe pipeline failed: %s\n",
                    error ? [[error localizedDescription] UTF8String] : "(unknown)");
            return 0;
        }
        return 1;
    }
#endif
    return 0;
}

static void ds4_gpu_detect_metal4_features(void) {
    g_metal4_runtime_available = 0;
    g_metal4_family_supported = 0;
    g_metal4_queue_supported = 0;
    g_metal4_m5_neural_accelerators_hint = 0;
    g_metal4_tensor_api_enabled = 0;
    g_metal4_tensor_api_compile_supported = 0;
    g_metal_device_name[0] = '\0';

    if (!g_device) return;

    const char *name = [[g_device name] UTF8String];
    if (name) {
        snprintf(g_metal_device_name, sizeof(g_metal_device_name), "%s", name);
    }

    const int metal4_disabled = ds4_gpu_env_bool("DS4_METAL_DISABLE_METAL4") > 0;

#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000
    if (@available(macOS 26.0, *)) {
        g_metal4_runtime_available = 1;
        g_metal4_family_supported =
            !metal4_disabled && [g_device supportsFamily:MTLGPUFamilyMetal4] ? 1 : 0;
        g_metal4_queue_supported = [g_device respondsToSelector:@selector(newMTL4CommandQueue)] ? 1 : 0;

        /*
         * Apple does not currently expose a separate "Neural Accelerator" bit
         * through Metal. On public M5 systems the hardware signal is the device
         * generation plus Metal 4 support, so keep this as a conservative hint.
         */
        if (g_metal4_family_supported && ds4_gpu_device_name_contains("M5")) {
            g_metal4_m5_neural_accelerators_hint = 1;
        }

        if (g_metal4_family_supported) {
            const int default_enable =
                ds4_gpu_device_name_contains("M5") ||
                ds4_gpu_device_name_contains("M6") ||
                ds4_gpu_device_name_contains("A19") ||
                ds4_gpu_device_name_contains("A20");

            /*
             * Metal 4 TensorOps are portable in source, but on pre-M5 hardware
             * they can map to ordinary shader fallbacks.  Keep the automatic
             * fast path restricted to hardware generations where the Neural
             * Accelerator/TensorOps path is expected to pay off; older Metal
             * machines continue to use the established kernels unless a future
             * device is explicitly added here.
             */
            if (default_enable) {
                g_metal4_tensor_api_compile_supported = ds4_gpu_compile_tensor_probe();
                g_metal4_tensor_api_enabled = g_metal4_tensor_api_compile_supported;
                if (!g_metal4_tensor_api_enabled) {
                    fprintf(stderr, "ds4: Metal 4 tensor API probe failed; using legacy Metal kernels\n");
                }
            } else {
                fprintf(stderr, "ds4: Metal 4 tensor API disabled for pre-M5/pre-A19 devices\n");
            }
        }
    }
#endif
}

static int ds4_gpu_warm_model_views(void) {
    if (g_model_view_count == 0) return 1;

    id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride");
    if (!pipeline) return 0;

    uint64_t stride = 1024ull * 1024ull;
    const char *stride_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_MB");
    if (stride_env && stride_env[0]) {
        char *end = NULL;
        unsigned long long mb = strtoull(stride_env, &end, 10);
        if (end != stride_env && mb > 0 && mb <= 1024) {
            stride = mb * 1024ull * 1024ull;
        }
    }
    const char *stride_kb_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_KB");
    if (stride_kb_env && stride_kb_env[0]) {
        char *end = NULL;
        unsigned long long kb = strtoull(stride_kb_env, &end, 10);
        if (end != stride_kb_env && kb > 0 && kb <= 1024ull * 1024ull) {
            stride = kb * 1024ull;
            const uint64_t page = (uint64_t)getpagesize();
            if (stride < page) stride = page;
        }
    }

    uint64_t total_touches = 0;
    for (uint32_t i = 0; i < g_model_view_count; i++) {
        total_touches += (g_model_views[i].bytes + stride - 1) / stride;
    }
    if (total_touches == 0 || total_touches > (uint64_t)NSUIntegerMax) return 0;

    const NSUInteger out_bytes = (NSUInteger)total_touches;
    id<MTLBuffer> out = [g_device newBufferWithLength:out_bytes
                                             options:MTLResourceStorageModeShared];
    if (!out) {
        fprintf(stderr, "ds4: Metal model warmup scratch allocation failed\n");
        return 0;
    }
    out.label = @"ds4_model_warmup";

    id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
    if (!cb) {
        fprintf(stderr, "ds4: Metal model warmup command buffer allocation failed\n");
        return 0;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    uint64_t dst_offset = 0;
    for (uint32_t i = 0; i < g_model_view_count; i++) {
        const uint64_t bytes = g_model_views[i].bytes;
        const uint64_t n = (bytes + stride - 1) / stride;
        [enc setBuffer:g_model_views[i].buffer offset:0 atIndex:0];
        [enc setBuffer:out offset:0 atIndex:1];
        [enc setBytes:&stride length:sizeof(stride) atIndex:2];
        [enc setBytes:&bytes length:sizeof(bytes) atIndex:3];
        [enc setBytes:&dst_offset length:sizeof(dst_offset) atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        dst_offset += n;
    }
    ds4_gpu_end_compute_encoder(cb, enc);

    [cb commit];
    [cb waitUntilCompleted];

    if (cb.status == MTLCommandBufferStatusError) {
        fprintf(stderr, "ds4: Metal model warmup failed: %s\n",
                [[cb.error localizedDescription] UTF8String]);
        return 0;
    }

    return 1;
}

static const char *ds4_gpu_mul_mm_id_map0_name(uint32_t ne20) {
    switch (ne20) {
        case 1:  return "kernel_mul_mm_id_map0_ne20_1";
        case 2:  return "kernel_mul_mm_id_map0_ne20_2";
        case 4:  return "kernel_mul_mm_id_map0_ne20_4";
        case 5:  return "kernel_mul_mm_id_map0_ne20_5";
        case 6:  return "kernel_mul_mm_id_map0_ne20_6";
        case 8:  return "kernel_mul_mm_id_map0_ne20_8";
        case 10: return "kernel_mul_mm_id_map0_ne20_10";
        case 16: return "kernel_mul_mm_id_map0_ne20_16";
        case 22: return "kernel_mul_mm_id_map0_ne20_22";
        default: return NULL;
    }
}

static id<MTLComputePipelineState> ds4_gpu_get_mul_mv_pipeline(
        const char *function_name,
        int16_t     nsg) {
    NSString *key = [NSString stringWithFormat:@"%s_nsg=%d", function_name, (int)nsg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600];

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_mul_mv_ext_pipeline(
        const char *function_name,
        int16_t     nsg,
        int16_t     nxpsg) {
    NSString *key = [NSString stringWithFormat:@"%s_nsg=%d_nxpsg=%d",
                     function_name, (int)nsg, (int)nxpsg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&nsg   type:MTLDataTypeShort atIndex:600];
    [constants setConstantValue:&nxpsg type:MTLDataTypeShort atIndex:601];

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_pad_pipeline(
        bool    has_mask,
        int32_t ncpsg) {
    NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_pad_mask=%d_ncpsg=%d",
                     has_mask ? 1 : 0, (int)ncpsg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:100];
    [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:125];

    NSError *error = nil;
    id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_pad"
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad function not found: %s\n",
                [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad pipeline failed: %s\n",
                [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_blk_pipeline(
        int32_t nqptg,
        int32_t ncpsg) {
    NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_blk_nqptg=%d_ncpsg=%d",
                     (int)nqptg, (int)ncpsg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&nqptg type:MTLDataTypeInt atIndex:224];
    [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:225];

    NSError *error = nil;
    id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_blk"
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk function not found: %s\n",
                [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk pipeline failed: %s\n",
                [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_pipeline(
        const char *function_name,
        bool        has_mask,
        bool        has_sinks,
        bool        has_bias,
        bool        has_scap,
        bool        has_kvpad,
        bool        bc_mask,
        int32_t     ns10,
        int32_t     ns20,
        int32_t     nsg) {
    NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d",
                     function_name,
                     has_mask ? 1 : 0,
                     has_sinks ? 1 : 0,
                     has_bias ? 1 : 0,
                     has_scap ? 1 : 0,
                     has_kvpad ? 1 : 0,
                     bc_mask ? 1 : 0,
                     (int)ns10,
                     (int)ns20,
                     (int)nsg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) return cached;

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&has_mask  type:MTLDataTypeBool atIndex:300];
    [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:301];
    [constants setConstantValue:&has_bias  type:MTLDataTypeBool atIndex:302];
    [constants setConstantValue:&has_scap  type:MTLDataTypeBool atIndex:303];
    [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:304];
    [constants setConstantValue:&bc_mask   type:MTLDataTypeBool atIndex:310];
    [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:320];
    [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:321];
    [constants setConstantValue:&nsg  type:MTLDataTypeInt atIndex:322];

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_vec_pipeline(
        const char *function_name,
        bool        has_mask,
        bool        has_sinks,
        bool        has_bias,
        bool        has_scap,
        bool        has_kvpad,
        bool        shared_kvpad,
        int32_t     ns10,
        int32_t     ns20,
        int32_t     nsg,
        int32_t     nwg) {
    /*
     * Decode calls this once per layer with identical arguments, so memoize
     * the last hit and skip the NSString key + dictionary lookup on the hot
     * path.  The generic cache below remains the fallback for new variants.
     */
    static struct {
        const char *fn;
        bool m, s, b, c, k, sp;
        int32_t n10, n20, sg, wg;
        id<MTLComputePipelineState> pipeline;
    } memo;
    if (memo.pipeline && memo.fn != NULL && strcmp(memo.fn, function_name) == 0 &&
        memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias &&
        memo.c == has_scap && memo.k == has_kvpad && memo.sp == shared_kvpad &&
        memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg && memo.wg == nwg) {
        return memo.pipeline;
    }

    NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_sharedpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
                     function_name,
                     has_mask ? 1 : 0,
                     has_sinks ? 1 : 0,
                     has_bias ? 1 : 0,
                     has_scap ? 1 : 0,
                     has_kvpad ? 1 : 0,
                     shared_kvpad ? 1 : 0,
                     (int)ns10,
                     (int)ns20,
                     (int)nsg,
                     (int)nwg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) {
        memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias,
                               has_scap, has_kvpad, shared_kvpad, ns10, ns20,
                               nsg, nwg, cached };
        return cached;
    }

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&has_mask  type:MTLDataTypeBool atIndex:400];
    [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:401];
    [constants setConstantValue:&has_bias  type:MTLDataTypeBool atIndex:402];
    [constants setConstantValue:&has_scap  type:MTLDataTypeBool atIndex:403];
    [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:404];
    [constants setConstantValue:&shared_kvpad type:MTLDataTypeBool atIndex:405];
    [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:420];
    [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:421];
    [constants setConstantValue:&nsg  type:MTLDataTypeInt atIndex:422];
    [constants setConstantValue:&nwg  type:MTLDataTypeInt atIndex:423];

    NSError *error = nil;
    NSString *name = [NSString stringWithUTF8String:function_name];
    id<MTLFunction> fn = [g_library newFunctionWithName:name
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal %s function not found: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
                function_name, [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias,
                           has_scap, has_kvpad, shared_kvpad, ns10, ns20,
                           nsg, nwg, pipeline };
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_reduce_pipeline(
        int32_t dv,
        int32_t nwg) {
    /* Same per-layer memo pattern as the vec getter above. */
    static int32_t memo_dv, memo_nwg;
    static id<MTLComputePipelineState> memo_pipeline;
    if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) {
        return memo_pipeline;
    }

    NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_dv=%d_nwg=%d",
                     (int)dv, (int)nwg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (cached) {
        memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached;
        return cached;
    }

    MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
    [constants setConstantValue:&dv  type:MTLDataTypeInt atIndex:500];
    [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501];

    NSError *error = nil;
    id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce"
                                         constantValues:constants
                                                  error:&error];
    if (!fn) {
        fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce function not found: %s\n",
                [[error localizedDescription] UTF8String]);
        return nil;
    }

    error = nil;
    id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
    if (!pipeline) {
        fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce pipeline failed: %s\n",
                [[error localizedDescription] UTF8String]);
        return nil;
    }

    [g_pipeline_cache setObject:pipeline forKey:key];
    memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline;
    return pipeline;
}

static uint32_t ds4_gpu_flash_attn_vec_nsg(uint32_t n_keys, uint32_t nwg, uint32_t ncpsg) {
    uint32_t nsg = 1;
    while (2u * nwg * nsg * ncpsg < n_keys && nsg < 4u) {
        nsg *= 2u;
    }
    return nsg;
}

static int ds4_gpu_trace_allocs(void) {
    static int initialized;
    static int enabled;
    if (!initialized) {
        enabled = getenv("DS4_METAL_TRACE_ALLOCS") != NULL;
        initialized = 1;
    }
    return enabled;
}

static double ds4_gpu_mib(uint64_t bytes) {
    return (double)bytes / (1024.0 * 1024.0);
}

static double ds4_gpu_gib(uint64_t bytes) {
    return (double)bytes / (1024.0 * 1024.0 * 1024.0);
}

static ds4_gpu_stream_expert_timing_snapshot
ds4_gpu_stream_expert_timing_current(void) {
    return (ds4_gpu_stream_expert_timing_snapshot) {
        .selected_calls = g_stream_expert_timing_selected_calls,
        .selected_read_ms = g_stream_expert_timing_selected_read_ms,
        .selected_sync_ms = g_stream_expert_timing_selected_sync_ms,
        .selected_copy_ms = g_stream_expert_timing_selected_copy_ms,
        .selected_bind_ms = g_stream_expert_timing_selected_bind_ms,
        .split_layers = g_stream_expert_timing_split_layers,
        .split_resident_experts = g_stream_expert_timing_split_resident_experts,
        .split_missing_experts = g_stream_expert_timing_split_missing_experts,
        .split_resident_ms = g_stream_expert_timing_split_resident_ms,
        .split_missing_ms = g_stream_expert_timing_split_missing_ms,
        .split_missing_load_ms =
            g_stream_expert_timing_split_missing_load_ms,
        .split_missing_slot_ms =
            g_stream_expert_timing_split_missing_slot_ms,
        .split_missing_prune_ms =
            g_stream_expert_timing_split_missing_prune_ms,
        .split_missing_addr_ms =
            g_stream_expert_timing_split_missing_addr_ms,
        .split_missing_wait_ms =
            g_stream_expert_timing_split_missing_wait_ms,
        .load_calls = g_stream_expert_timing_load_calls,
        .load_prepare_ms = g_stream_expert_timing_load_prepare_ms,
        .load_pread_ms = g_stream_expert_timing_load_pread_ms,
        .load_modify_ms = g_stream_expert_timing_load_modify_ms,
        .load_install_ms = g_stream_expert_timing_load_install_ms,
        .prepare_batch_reuse_calls =
            g_stream_expert_timing_prepare_batch_reuse_calls,
        .prepare_batch_reuse_ms =
            g_stream_expert_timing_prepare_batch_reuse_ms,
        .prepare_buffer_calls =
            g_stream_expert_timing_prepare_buffer_calls,
        .prepare_buffer_ms =
            g_stream_expert_timing_prepare_buffer_ms,
        .prepare_task_experts =
            g_stream_expert_timing_prepare_task_experts,
        .prepare_task_ms =
            g_stream_expert_timing_prepare_task_ms,
        .reuse_scan_calls =
            g_stream_expert_timing_reuse_scan_calls,
        .reuse_scan_entries =
            g_stream_expert_timing_reuse_scan_entries,
        .reuse_scan_ms =
            g_stream_expert_timing_reuse_scan_ms,
        .reuse_clear_ms =
            g_stream_expert_timing_reuse_clear_ms,
        .readahead_calls =
            g_stream_expert_timing_readahead_calls,
        .readahead_bytes =
            g_stream_expert_timing_readahead_bytes,
        .readahead_ms =
            g_stream_expert_timing_readahead_ms,
        .cache_all_resident_layers =
            g_stream_expert_timing_cache_all_resident_layers,
        .cache_all_missing_layers =
            g_stream_expert_timing_cache_all_missing_layers,
        .cache_mixed_layers = g_stream_expert_timing_cache_mixed_layers,
        .cache_resident_experts =
            g_stream_expert_timing_cache_resident_experts,
        .cache_missing_experts =
            g_stream_expert_timing_cache_missing_experts,
    };
}

static uint64_t ds4_gpu_stream_expert_timing_delta_u64(
        uint64_t current,
        uint64_t previous) {
    return current >= previous ? current - previous : current;
}

static double ds4_gpu_stream_expert_timing_delta_f64(
        double current,
        double previous) {
    return current >= previous ? current - previous : current;
}

static ds4_gpu_stream_expert_timing_snapshot
ds4_gpu_stream_expert_timing_delta(
        ds4_gpu_stream_expert_timing_snapshot current,
        ds4_gpu_stream_expert_timing_snapshot previous) {
    return (ds4_gpu_stream_expert_timing_snapshot) {
        .selected_calls =
            ds4_gpu_stream_expert_timing_delta_u64(current.selected_calls,
                                                   previous.selected_calls),
        .selected_read_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.selected_read_ms,
                                                   previous.selected_read_ms),
        .selected_sync_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.selected_sync_ms,
                                                   previous.selected_sync_ms),
        .selected_copy_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.selected_copy_ms,
                                                   previous.selected_copy_ms),
        .selected_bind_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.selected_bind_ms,
                                                   previous.selected_bind_ms),
        .split_layers =
            ds4_gpu_stream_expert_timing_delta_u64(current.split_layers,
                                                   previous.split_layers),
        .split_resident_experts =
            ds4_gpu_stream_expert_timing_delta_u64(current.split_resident_experts,
                                                   previous.split_resident_experts),
        .split_missing_experts =
            ds4_gpu_stream_expert_timing_delta_u64(current.split_missing_experts,
                                                   previous.split_missing_experts),
        .split_resident_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.split_resident_ms,
                                                   previous.split_resident_ms),
        .split_missing_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.split_missing_ms,
                                                   previous.split_missing_ms),
        .split_missing_load_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.split_missing_load_ms,
                    previous.split_missing_load_ms),
        .split_missing_slot_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.split_missing_slot_ms,
                    previous.split_missing_slot_ms),
        .split_missing_prune_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.split_missing_prune_ms,
                    previous.split_missing_prune_ms),
        .split_missing_addr_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.split_missing_addr_ms,
                    previous.split_missing_addr_ms),
        .split_missing_wait_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.split_missing_wait_ms,
                    previous.split_missing_wait_ms),
        .load_calls =
            ds4_gpu_stream_expert_timing_delta_u64(current.load_calls,
                                                   previous.load_calls),
        .load_prepare_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.load_prepare_ms,
                                                   previous.load_prepare_ms),
        .load_pread_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.load_pread_ms,
                                                   previous.load_pread_ms),
        .load_modify_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.load_modify_ms,
                                                   previous.load_modify_ms),
        .load_install_ms =
            ds4_gpu_stream_expert_timing_delta_f64(current.load_install_ms,
                                                   previous.load_install_ms),
        .prepare_batch_reuse_calls =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.prepare_batch_reuse_calls,
                    previous.prepare_batch_reuse_calls),
        .prepare_batch_reuse_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.prepare_batch_reuse_ms,
                    previous.prepare_batch_reuse_ms),
        .prepare_buffer_calls =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.prepare_buffer_calls,
                    previous.prepare_buffer_calls),
        .prepare_buffer_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.prepare_buffer_ms,
                    previous.prepare_buffer_ms),
        .prepare_task_experts =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.prepare_task_experts,
                    previous.prepare_task_experts),
        .prepare_task_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.prepare_task_ms,
                    previous.prepare_task_ms),
        .reuse_scan_calls =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.reuse_scan_calls,
                    previous.reuse_scan_calls),
        .reuse_scan_entries =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.reuse_scan_entries,
                    previous.reuse_scan_entries),
        .reuse_scan_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.reuse_scan_ms,
                    previous.reuse_scan_ms),
        .reuse_clear_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.reuse_clear_ms,
                    previous.reuse_clear_ms),
        .readahead_calls =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.readahead_calls,
                    previous.readahead_calls),
        .readahead_bytes =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.readahead_bytes,
                    previous.readahead_bytes),
        .readahead_ms =
            ds4_gpu_stream_expert_timing_delta_f64(
                    current.readahead_ms,
                    previous.readahead_ms),
        .cache_all_resident_layers =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.cache_all_resident_layers,
                    previous.cache_all_resident_layers),
        .cache_all_missing_layers =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.cache_all_missing_layers,
                    previous.cache_all_missing_layers),
        .cache_mixed_layers =
            ds4_gpu_stream_expert_timing_delta_u64(current.cache_mixed_layers,
                                                   previous.cache_mixed_layers),
        .cache_resident_experts =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.cache_resident_experts,
                    previous.cache_resident_experts),
        .cache_missing_experts =
            ds4_gpu_stream_expert_timing_delta_u64(
                    current.cache_missing_experts,
                    previous.cache_missing_experts),
    };
}

static int ds4_gpu_stream_expert_timing_has_data(
        ds4_gpu_stream_expert_timing_snapshot s) {
    return s.selected_calls != 0 ||
           s.split_layers != 0 ||
           s.load_calls != 0 ||
           s.cache_all_resident_layers != 0 ||
           s.cache_all_missing_layers != 0 ||
           s.cache_mixed_layers != 0;
}

static void ds4_gpu_stream_expert_timing_print(
        const char *scope,
        ds4_gpu_stream_expert_timing_snapshot s) {
    if (!ds4_gpu_stream_expert_timing_has_data(s)) return;
    const double selected_calls = (double)s.selected_calls;
    const double split_layers = (double)s.split_layers;
    const double selected_read_avg =
        selected_calls != 0.0 ? s.selected_read_ms / selected_calls : 0.0;
    const double selected_sync_avg =
        selected_calls != 0.0 ? s.selected_sync_ms / selected_calls : 0.0;
    const double selected_copy_avg =
        selected_calls != 0.0 ? s.selected_copy_ms / selected_calls : 0.0;
    const double selected_bind_avg =
        selected_calls != 0.0 ? s.selected_bind_ms / selected_calls : 0.0;
    const double split_resident_avg =
        split_layers != 0.0 ? s.split_resident_ms / split_layers : 0.0;
    const double split_missing_avg =
        split_layers != 0.0 ? s.split_missing_ms / split_layers : 0.0;
    const double split_missing_load_avg =
        split_layers != 0.0 ?
            s.split_missing_load_ms / split_layers : 0.0;
    const double split_missing_slot_avg =
        split_layers != 0.0 ?
            s.split_missing_slot_ms / split_layers : 0.0;
    const double split_missing_prune_avg =
        split_layers != 0.0 ?
            s.split_missing_prune_ms / split_layers : 0.0;
    const double split_missing_addr_avg =
        split_layers != 0.0 ?
            s.split_missing_addr_ms / split_layers : 0.0;
    const double split_missing_wait_avg =
        split_layers != 0.0 ?
            s.split_missing_wait_ms / split_layers : 0.0;
    const double load_calls = (double)s.load_calls;
    const double load_prepare_avg =
        load_calls != 0.0 ? s.load_prepare_ms / load_calls : 0.0;
    const double load_pread_avg =
        load_calls != 0.0 ? s.load_pread_ms / load_calls : 0.0;
    const double load_modify_avg =
        load_calls != 0.0 ? s.load_modify_ms / load_calls : 0.0;
    const double load_install_avg =
        load_calls != 0.0 ? s.load_install_ms / load_calls : 0.0;
    const double prepare_batch_reuse_avg =
        s.prepare_batch_reuse_calls != 0 ?
            s.prepare_batch_reuse_ms /
                (double)s.prepare_batch_reuse_calls : 0.0;
    const double prepare_buffer_avg =
        s.prepare_buffer_calls != 0 ?
            s.prepare_buffer_ms / (double)s.prepare_buffer_calls : 0.0;
    const double prepare_task_avg =
        s.prepare_task_experts != 0 ?
            s.prepare_task_ms / (double)s.prepare_task_experts : 0.0;
    const double reuse_scan_avg =
        s.reuse_scan_calls != 0 ?
            s.reuse_scan_ms / (double)s.reuse_scan_calls : 0.0;
    const double reuse_scan_entries_avg =
        s.reuse_scan_calls != 0 ?
            (double)s.reuse_scan_entries / (double)s.reuse_scan_calls : 0.0;
    const double readahead_avg =
        s.readahead_calls != 0 ?
            s.readahead_ms / (double)s.readahead_calls : 0.0;
    const double split_resident_experts_avg =
        split_layers != 0.0 ?
            (double)s.split_resident_experts / split_layers : 0.0;
    const double split_missing_experts_avg =
        split_layers != 0.0 ?
            (double)s.split_missing_experts / split_layers : 0.0;
    const uint64_t cache_layers =
        s.cache_all_resident_layers +
        s.cache_all_missing_layers +
        s.cache_mixed_layers;
    const double cache_layer_count = (double)cache_layers;
    const double cache_resident_experts_avg =
        cache_layer_count != 0.0 ?
            (double)s.cache_resident_experts / cache_layer_count : 0.0;
    const double cache_missing_experts_avg =
        cache_layer_count != 0.0 ?
            (double)s.cache_missing_experts / cache_layer_count : 0.0;
    fprintf(stderr,
            "ds4:   streaming expert timing %s selected_calls=%llu read_avg=%.3f ms sync_avg=%.3f ms copy_avg=%.3f ms bind_avg=%.3f ms read_total=%.3f ms sync_total=%.3f ms copy_total=%.3f ms bind_total=%.3f ms split_layers=%llu resident_experts_avg=%.2f missing_experts_avg=%.2f resident_submit_avg=%.3f ms missing_bind_avg=%.3f ms resident_submit_total=%.3f ms missing_bind_total=%.3f ms missing_load_avg=%.3f ms missing_slot_avg=%.3f ms missing_prune_avg=%.3f ms missing_addr_avg=%.3f ms missing_wait_avg=%.3f ms missing_wait_total=%.3f ms load_calls=%llu load_prepare_avg=%.3f ms load_pread_avg=%.3f ms load_modify_avg=%.3f ms load_install_avg=%.3f ms prepare_batch_reuse_calls=%llu prepare_batch_reuse_avg=%.3f ms prepare_batch_reuse_total=%.3f ms prepare_buffer_calls=%llu prepare_buffer_avg=%.3f ms prepare_buffer_total=%.3f ms prepare_task_experts=%llu prepare_task_avg=%.3f ms prepare_task_total=%.3f ms reuse_scan_calls=%llu reuse_scan_entries_avg=%.1f reuse_scan_avg=%.3f ms reuse_scan_total=%.3f ms reuse_clear_total=%.3f ms readahead_calls=%llu readahead_avg=%.3f ms readahead_total=%.3f ms readahead_gib=%.2f cache_all_resident=%llu cache_all_missing=%llu cache_mixed=%llu cache_resident_avg=%.2f cache_missing_avg=%.2f\n",
            scope ? scope : "total",
            (unsigned long long)s.selected_calls,
            selected_read_avg,
            selected_sync_avg,
            selected_copy_avg,
            selected_bind_avg,
            s.selected_read_ms,
            s.selected_sync_ms,
            s.selected_copy_ms,
            s.selected_bind_ms,
            (unsigned long long)s.split_layers,
            split_resident_experts_avg,
            split_missing_experts_avg,
            split_resident_avg,
            split_missing_avg,
            s.split_resident_ms,
            s.split_missing_ms,
            split_missing_load_avg,
            split_missing_slot_avg,
            split_missing_prune_avg,
            split_missing_addr_avg,
            split_missing_wait_avg,
            s.split_missing_wait_ms,
            (unsigned long long)s.load_calls,
            load_prepare_avg,
            load_pread_avg,
            load_modify_avg,
            load_install_avg,
            (unsigned long long)s.prepare_batch_reuse_calls,
            prepare_batch_reuse_avg,
            s.prepare_batch_reuse_ms,
            (unsigned long long)s.prepare_buffer_calls,
            prepare_buffer_avg,
            s.prepare_buffer_ms,
            (unsigned long long)s.prepare_task_experts,
            prepare_task_avg,
            s.prepare_task_ms,
            (unsigned long long)s.reuse_scan_calls,
            reuse_scan_entries_avg,
            reuse_scan_avg,
            s.reuse_scan_ms,
            s.reuse_clear_ms,
            (unsigned long long)s.readahead_calls,
            readahead_avg,
            s.readahead_ms,
            ds4_gpu_gib(s.readahead_bytes),
            (unsigned long long)s.cache_all_resident_layers,
            (unsigned long long)s.cache_all_missing_layers,
            (unsigned long long)s.cache_mixed_layers,
            cache_resident_experts_avg,
            cache_missing_experts_avg);
}

static void ds4_gpu_print_task_memory_report(void) {
    task_vm_info_data_t info;
    mach_msg_type_number_t count = TASK_VM_INFO_COUNT;
    const kern_return_t kr = task_info(mach_task_self(),
                                       TASK_VM_INFO,
                                       (task_info_t)&info,
                                       &count);
    if (kr != KERN_SUCCESS) return;

    fprintf(stderr,
            "ds4:   macOS task memory footprint %.2f GiB, resident %.2f GiB, virtual %.2f GiB\n",
            ds4_gpu_gib((uint64_t)info.phys_footprint),
            ds4_gpu_gib((uint64_t)info.resident_size),
            ds4_gpu_gib((uint64_t)info.virtual_size));
}

void ds4_gpu_print_memory_report(const char *label) {
    uint64_t cached_prefill_mask_bytes = 0;
    uint64_t cached_prefill_blk_bytes = 0;
    for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) {
        const ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry =
            &g_zero_prefix_prefill_mask_cache[i];
        if (entry->mask) cached_prefill_mask_bytes += entry->mask_bytes;
        if (entry->blk) cached_prefill_blk_bytes += entry->blk_bytes;
    }
    const uint64_t scratch =
        (uint64_t)g_flash_attn_mask_bytes +
        (uint64_t)g_flash_attn_zero_mask_bytes +
        cached_prefill_mask_bytes +
        (uint64_t)g_flash_attn_pad_bytes +
        (uint64_t)g_flash_attn_tmp_bytes +
        (uint64_t)g_flash_attn_blk_bytes +
        cached_prefill_blk_bytes +
        (uint64_t)g_flash_attn_ring_bytes +
        (uint64_t)g_flash_attn_kv_bytes +
        (uint64_t)g_glm_flash_attn_mask_bytes +
        (uint64_t)g_compressor_pool_kv_bytes +
        (uint64_t)g_compressor_pool_score_bytes +
        (uint64_t)g_compressor_pool_score_cont_bytes +
        (uint64_t)g_compressor_pool_softmax_bytes +
        (uint64_t)g_compressor_pool_product_bytes +
        (uint64_t)g_compressor_store_ape_bytes +
        (uint64_t)g_compressor_store_score_bytes +
        (uint64_t)g_embed_rows_bytes +
        (uint64_t)g_router_selection_bytes +
        (uint64_t)g_router_weight_sum_bytes +
        (uint64_t)g_indexer_head_scores_bytes +
        (uint64_t)g_indexer_topk_bytes +
        (uint64_t)g_indexed_topk_bytes +
        (uint64_t)g_f16_round_scratch_bytes +
        (uint64_t)g_raw_store_round_bytes +
        (uint64_t)g_moe_gate_scratch_bytes +
        (uint64_t)g_moe_down_scratch_bytes +
        (uint64_t)g_moe_id_map_bytes +
        (uint64_t)g_moe_q4_gate_slots_bytes +
        (uint64_t)g_moe_q4_up_slots_bytes +
        (uint64_t)g_moe_q4_down_slots_bytes;

    pthread_mutex_lock(&g_tensor_mu);
    const uint64_t tensor_live_snap = g_tensor_alloc_live_bytes;
    const uint64_t tensor_peak_snap = g_tensor_alloc_peak_bytes;
    pthread_mutex_unlock(&g_tensor_mu);

    uint64_t tracked_live = tensor_live_snap;
    if (tracked_live > UINT64_MAX - g_stream_expert_cache_bytes) {
        tracked_live = UINT64_MAX;
    } else {
        tracked_live += g_stream_expert_cache_bytes;
    }

    const bool color = ds4_log_is_tty(stderr);
    const char *green = color ? "\x1b[32m" : "";
    const char *bright_green = color ? "\x1b[1;32m" : "";
    const char *reset = color ? "\x1b[0m" : "";
    fprintf(stderr,
            "%sds4: Metal memory%s%s: runtime %.2f GiB + streaming experts %.2f GiB = %s%.2f GiB tracked live%s\n",
            green,
            label && label[0] ? " " : "",
            label && label[0] ? label : "",
            ds4_gpu_gib(tensor_live_snap),
            ds4_gpu_gib(g_stream_expert_cache_bytes),
            bright_green,
            ds4_gpu_gib(tracked_live),
            reset);
    if (color) fputs(green, stderr);
    fprintf(stderr,
            "ds4:   runtime tensors live %.2f MiB peak %.2f MiB\n",
            ds4_gpu_mib(tensor_live_snap),
            ds4_gpu_mib(tensor_peak_snap));
    ds4_gpu_print_task_memory_report();
    fprintf(stderr,
            "ds4:   mmap model wrapper spans %llu buffers %.2f GiB total, %.2f GiB max (not copied)\n",
            (unsigned long long)g_model_wrap_count,
            ds4_gpu_gib(g_model_wrap_bytes),
            ds4_gpu_gib(g_model_wrap_max_bytes));
    if (g_model_buffer_cache && [g_model_buffer_cache count] != 0) {
        const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes();
        if (limit == 0) {
            fprintf(stderr,
                    "ds4:   exact model view cache %lu buffers %.2f GiB unlimited, %llu evictions (not copied)\n",
                    (unsigned long)[g_model_buffer_cache count],
                    ds4_gpu_gib(g_model_buffer_cache_bytes),
                    (unsigned long long)g_model_buffer_cache_evictions);
        } else {
            fprintf(stderr,
                    "ds4:   exact model view cache %lu buffers %.2f GiB / %.2f GiB, %llu evictions (not copied)\n",
                    (unsigned long)[g_model_buffer_cache count],
                    ds4_gpu_gib(g_model_buffer_cache_bytes),
                    ds4_gpu_gib(limit),
                    (unsigned long long)g_model_buffer_cache_evictions);
        }
    }
    if (g_stream_expert_cache_hits != 0 ||
        g_stream_expert_cache_misses != 0 ||
        g_stream_expert_cache_bytes != 0) {
        const uint64_t budget = ds4_gpu_stream_expert_cache_configured_budget();
        uint64_t target_bytes = 0;
        if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) {
            target_bytes =
                budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ?
                    UINT64_MAX :
                    budget * g_stream_expert_cache_expert_bytes;
        }
        const uint64_t lookups = g_stream_expert_cache_hits + g_stream_expert_cache_misses;
        const double hit_rate = lookups ?
            (double)g_stream_expert_cache_hits / (double)lookups : 0.0;
        if (g_stream_expert_cache_evict_advise_bytes != 0 ||
            g_stream_expert_cache_willneed_advise_bytes != 0 ||
            g_stream_expert_cache_pread_bytes != 0) {
            fprintf(stderr,
                    "ds4:   streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu evict_dontneed=%.2f GiB miss_willneed=%.2f GiB miss_pread=%.2f GiB pread_ms=%.3f\n",
                    (unsigned long long)budget,
                    g_stream_expert_cache_entry_count,
                    ds4_gpu_mib(g_stream_expert_cache_expert_bytes),
                    ds4_gpu_gib(target_bytes),
                    ds4_gpu_gib(g_stream_expert_cache_bytes),
                    (unsigned long long)g_stream_expert_cache_hits,
                    (unsigned long long)g_stream_expert_cache_misses,
                    hit_rate,
                    (unsigned long long)g_stream_expert_cache_wraps,
                    (unsigned long long)g_stream_expert_cache_evictions,
                    (unsigned long long)g_stream_expert_cache_buffer_allocs,
                    (unsigned long long)g_stream_expert_cache_buffer_reuses,
                    ds4_gpu_gib(g_stream_expert_cache_evict_advise_bytes),
                    ds4_gpu_gib(g_stream_expert_cache_willneed_advise_bytes),
                    ds4_gpu_gib(g_stream_expert_cache_pread_bytes),
                    g_stream_expert_cache_pread_ms);
        } else {
            fprintf(stderr,
                    "ds4:   streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu\n",
                    (unsigned long long)budget,
                    g_stream_expert_cache_entry_count,
                    ds4_gpu_mib(g_stream_expert_cache_expert_bytes),
                    ds4_gpu_gib(target_bytes),
                    ds4_gpu_gib(g_stream_expert_cache_bytes),
                    (unsigned long long)g_stream_expert_cache_hits,
                    (unsigned long long)g_stream_expert_cache_misses,
                    hit_rate,
                    (unsigned long long)g_stream_expert_cache_wraps,
                    (unsigned long long)g_stream_expert_cache_evictions,
                    (unsigned long long)g_stream_expert_cache_buffer_allocs,
                    (unsigned long long)g_stream_expert_cache_buffer_reuses);
        }
        if (ds4_gpu_stream_expert_timing_summary_enabled()) {
            const ds4_gpu_stream_expert_timing_snapshot total =
                ds4_gpu_stream_expert_timing_current();
            if (ds4_gpu_stream_expert_timing_has_data(total)) {
                const ds4_gpu_stream_expert_timing_snapshot delta =
                    ds4_gpu_stream_expert_timing_delta(
                            total,
                            g_stream_expert_timing_last_report);
                ds4_gpu_stream_expert_timing_print("total", total);
                ds4_gpu_stream_expert_timing_print("delta", delta);
                g_stream_expert_timing_last_report = total;
            }
        }
        if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS") != NULL) {
            fprintf(stderr, "ds4:   streaming expert cache per-layer stats:\n");
            for (uint32_t layer = 0;
                 layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
                 layer++) {
                const uint64_t hits = g_stream_expert_cache_layer_hits[layer];
                const uint64_t misses = g_stream_expert_cache_layer_misses[layer];
                const uint64_t lookups = hits + misses;
                const uint64_t evictions =
                    g_stream_expert_cache_layer_evictions[layer];
                const uint64_t pread_bytes =
                    g_stream_expert_cache_layer_pread_bytes[layer];
                const double pread_ms =
                    g_stream_expert_cache_layer_pread_ms[layer];
                const uint32_t cached =
                    g_stream_expert_cache_layer_count[layer];
                const uint32_t layer_slots =
                    ds4_gpu_stream_expert_cache_configured_count() != 0 ?
                        DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT : 0;
                if (lookups == 0 && evictions == 0 && pread_bytes == 0 &&
                    cached == 0) {
                    continue;
                }
                const double layer_hit_rate = lookups ?
                    (double)hits / (double)lookups : 0.0;
                fprintf(stderr,
                        "ds4:     layer=%u layer_slots=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n",
                        layer,
                        layer_slots,
                        cached,
                        (unsigned long long)hits,
                        (unsigned long long)misses,
                        layer_hit_rate,
                        (unsigned long long)evictions,
                        ds4_gpu_gib(pread_bytes),
                        pread_ms);
            }
            if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS_DELTA") != NULL) {
                fprintf(stderr, "ds4:   streaming expert cache per-layer delta:\n");
                for (uint32_t layer = 0;
                     layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
                     layer++) {
                    const uint64_t hits =
                        g_stream_expert_cache_layer_hits[layer];
                    const uint64_t misses =
                        g_stream_expert_cache_layer_misses[layer];
                    const uint64_t evictions =
                        g_stream_expert_cache_layer_evictions[layer];
                    const uint64_t pread_bytes =
                        g_stream_expert_cache_layer_pread_bytes[layer];
                    const double pread_ms =
                        g_stream_expert_cache_layer_pread_ms[layer];
                    const uint64_t delta_hits =
                        hits - g_stream_expert_cache_layer_last_hits[layer];
                    const uint64_t delta_misses =
                        misses - g_stream_expert_cache_layer_last_misses[layer];
                    const uint64_t delta_evictions =
                        evictions -
                        g_stream_expert_cache_layer_last_evictions[layer];
                    const uint64_t delta_pread_bytes =
                        pread_bytes -
                        g_stream_expert_cache_layer_last_pread_bytes[layer];
                    const double delta_pread_ms =
                        pread_ms -
                        g_stream_expert_cache_layer_last_pread_ms[layer];
                    const uint64_t lookups = delta_hits + delta_misses;
                    if (lookups == 0 && delta_evictions == 0 &&
                        delta_pread_bytes == 0) {
                        continue;
                    }
                    const double hit_rate = lookups ?
                        (double)delta_hits / (double)lookups : 0.0;
                    fprintf(stderr,
                            "ds4:     layer=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n",
                            layer,
                            g_stream_expert_cache_layer_count[layer],
                            (unsigned long long)delta_hits,
                            (unsigned long long)delta_misses,
                            hit_rate,
                            (unsigned long long)delta_evictions,
                            ds4_gpu_gib(delta_pread_bytes),
                            delta_pread_ms);
                    g_stream_expert_cache_layer_last_hits[layer] = hits;
                    g_stream_expert_cache_layer_last_misses[layer] = misses;
                    g_stream_expert_cache_layer_last_evictions[layer] =
                        evictions;
                    g_stream_expert_cache_layer_last_pread_bytes[layer] =
                        pread_bytes;
                    g_stream_expert_cache_layer_last_pread_ms[layer] =
                        pread_ms;
                }
            }
        }
    }
    fprintf(stderr,
            "ds4:   model residency requests %llu%s\n",
            (unsigned long long)g_model_residency_count,
            g_ssd_streaming_mode ? " (ssd-streaming)" :
            (getenv("DS4_METAL_NO_RESIDENCY") != NULL ? " (disabled)" : ""));
    fprintf(stderr,
            "ds4:   device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n",
            g_metal_device_name[0] ? g_metal_device_name : "(unknown)",
            g_metal4_runtime_available ? "yes" : "no",
            g_metal4_family_supported ? "yes" : "no",
            g_metal4_queue_supported ? "yes" : "no",
            g_metal4_tensor_api_enabled ? "enabled" :
                (g_metal4_tensor_api_compile_supported ? "available" : "disabled"),
            g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected");
    fprintf(stderr,
            "ds4:   accelerated Metal path %s%s\n",
            ds4_gpu_mpp_available() ? "enabled" : "disabled",
            g_quality_mode ? " by --quality" :
                (!g_metal4_tensor_api_enabled ? " (tensor API unavailable)" : ""));
    fprintf(stderr,
            "ds4:   device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n",
            g_metal_device_name[0] ? g_metal_device_name : "(unknown)",
            g_metal4_runtime_available ? "yes" : "no",
            g_metal4_family_supported ? "yes" : "no",
            g_metal4_queue_supported ? "yes" : "no",
            g_metal4_tensor_api_enabled ? "enabled" :
                (g_metal4_tensor_api_compile_supported ? "available" : "disabled"),
            g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected");
    fprintf(stderr,
            "ds4:   scratch %.2f MiB (flash mask %.2f, pad %.2f, tmp %.2f, blk %.2f, ring %.2f, kv %.2f, compressor %.2f, router %.2f, indexer %.2f, moe %.2f, f16 %.2f, raw-store %.2f)\n",
            ds4_gpu_mib(scratch),
            ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes +
                        (uint64_t)g_glm_flash_attn_mask_bytes +
                        (uint64_t)g_flash_attn_zero_mask_bytes +
                        cached_prefill_mask_bytes),
            ds4_gpu_mib((uint64_t)g_flash_attn_pad_bytes),
            ds4_gpu_mib((uint64_t)g_flash_attn_tmp_bytes),
            ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes +
                        cached_prefill_blk_bytes),
            ds4_gpu_mib((uint64_t)g_flash_attn_ring_bytes),
            ds4_gpu_mib((uint64_t)g_flash_attn_kv_bytes),
            ds4_gpu_mib((uint64_t)g_compressor_pool_kv_bytes +
                          (uint64_t)g_compressor_pool_score_bytes +
                          (uint64_t)g_compressor_pool_score_cont_bytes +
                          (uint64_t)g_compressor_pool_softmax_bytes +
                          (uint64_t)g_compressor_pool_product_bytes +
                          (uint64_t)g_compressor_store_ape_bytes +
                          (uint64_t)g_compressor_store_score_bytes +
                          (uint64_t)g_embed_rows_bytes),
            ds4_gpu_mib((uint64_t)g_router_selection_bytes +
                          (uint64_t)g_router_weight_sum_bytes),
            ds4_gpu_mib((uint64_t)g_indexer_head_scores_bytes +
                          (uint64_t)g_indexer_topk_bytes +
                          (uint64_t)g_indexed_topk_bytes),
            ds4_gpu_mib((uint64_t)g_moe_gate_scratch_bytes +
                          (uint64_t)g_moe_down_scratch_bytes +
                          (uint64_t)g_moe_id_map_bytes +
                          (uint64_t)g_moe_q4_gate_slots_bytes +
                          (uint64_t)g_moe_q4_up_slots_bytes +
                          (uint64_t)g_moe_q4_down_slots_bytes),
            ds4_gpu_mib((uint64_t)g_f16_round_scratch_bytes),
            ds4_gpu_mib((uint64_t)g_raw_store_round_bytes));
    if (color) fputs(reset, stderr);
}

void ds4_gpu_set_quality(bool quality) {
    g_quality_mode = quality ? 1 : 0;
}

void ds4_gpu_set_glm_model(bool enabled) {
    g_glm_model_mode = enabled ? 1 : 0;
}

void ds4_gpu_set_ssd_streaming(bool enabled) {
    g_ssd_streaming_mode = enabled ? 1 : 0;
    ds4_gpu_stream_expert_cache_clear_all(1);
    if (g_ssd_streaming_mode) {
        fprintf(stderr,
                "ds4: Metal SSD streaming mode enabled; full model residency and warmup are skipped\n");
    }
}

void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) {
    g_glm_streaming_prefill_full_layer_runtime = enabled ? 1 : 0;
}

void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) {
    if (experts > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
        experts = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES;
    }
    g_stream_expert_cache_budget_override = experts;
    ds4_gpu_stream_expert_cache_clear_all(1);
}

void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) {
    /*
     * Pre-seed the cache's single slab size class with the model's uniform
     * per-expert bytes (first routed layer). With a mixed-precision GGUF this
     * pins the class to the majority layers so the boosted ones are rejected
     * deterministically from startup, instead of depending on which layer
     * happens to touch the cache first.
     */
    g_stream_expert_cache_expert_bytes = bytes;
}

uint64_t ds4_gpu_recommended_working_set_size(void) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!g_device) return 0;
    return (uint64_t)[g_device recommendedMaxWorkingSetSize];
}

static int ds4_gpu_model_map_log_enabled(void) {
    if (!g_ssd_streaming_mode) return 1;
    const char *trace = getenv("DS4_METAL_STREAMING_MAP_TRACE");
    return trace && trace[0] && strcmp(trace, "0") != 0;
}

static id<MTLBuffer> ds4_gpu_wrap_model_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset);

static id<MTLBuffer> ds4_gpu_wrap_model_exact_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset);

static const char *ds4_gpu_source =
"#include <metal_stdlib>\n"
"#ifdef DS4_METAL_HAS_TENSOR\n"
"#include <metal_tensor>\n"
"#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>\n"
"#endif\n"
"using namespace metal;\n"
"#ifdef DS4_METAL_HAS_TENSOR\n"
"using namespace mpp::tensor_ops;\n"
"#endif\n"
"\n"
"#define MAX(x, y) ((x) > (y) ? (x) : (y))\n"
"#define MIN(x, y) ((x) < (y) ? (x) : (y))\n"
"#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }\n"
"#define QK8_0 32\n"
"#ifndef QK_K\n"
"#define QK_K 256\n"
"#endif\n"
"#define N_SIMDWIDTH 32\n"
"#define N_R0_Q8_0 2\n"
"#define N_SG_Q8_0 4\n"
"#define FC_MUL_MV 600\n"
"#define FC_MUL_MM 700\n"
"#define FC_BIN 1300\n"
"#define FOR_UNROLL(x) _Pragma(\"clang loop unroll(full)\") for (x)\n"
"#define M_PI_F 3.14159265358979323846f\n"
"\n"
"// Reads one byte per stride to warm model-backed pages without copying the\n"
"// model. This is outside inference and exists only to reduce first-use stalls.\n"
"kernel void kernel_touch_u8_stride(\n"
"        device const uchar    *src        [[buffer(0)]],\n"
"        device uchar          *dst        [[buffer(1)]],\n"
"        constant ulong        &stride     [[buffer(2)]],\n"
"        constant ulong        &bytes      [[buffer(3)]],\n"
"        constant ulong        &dst_offset [[buffer(4)]],\n"
"        uint gid [[thread_position_in_grid]]) {\n"
"    ulong off = (ulong)gid * stride;\n"
"    if (off >= bytes) return;\n"
"    dst[dst_offset + (ulong)gid] = src[off];\n"
"}\n"
"\n"
"enum ds4_sort_order {\n"
"    DS4_SORT_ORDER_ASC,\n"
"    DS4_SORT_ORDER_DESC,\n"
"};\n"
"\n"
"struct block_q8_0 {\n"
"    half d;\n"
"    int8_t qs[QK8_0];\n"
"};\n"
"\n"
"struct block_q8_K {\n"
"    float d;\n"
"    int8_t qs[QK_K];\n"
"    int16_t bsums[QK_K / 16];\n"
"};\n"
"\n"
"\n";

static NSString *ds4_gpu_full_source(void) {
    NSString *base = [NSString stringWithUTF8String:ds4_gpu_source];
    NSFileManager *fm = [NSFileManager defaultManager];
    /*
     * Kernels are kept as separate files for review, then concatenated into one
     * Metal library.  Environment overrides are still honored so a diagnostic
     * run can swap one source file without changing the executable.
     */
    NSArray<NSArray<NSString *> *> *required_sources = @[
        @[@"DS4_METAL_FLASH_ATTN_SOURCE", @"metal/flash_attn.metal"],
        @[@"DS4_METAL_DENSE_SOURCE",      @"metal/dense.metal"],
        @[@"DS4_METAL_MOE_SOURCE",        @"metal/moe.metal"],
        @[@"DS4_METAL_DSV4_HC_SOURCE",    @"metal/dsv4_hc.metal"],
        @[@"DS4_METAL_UNARY_SOURCE",      @"metal/unary.metal"],
        @[@"DS4_METAL_DSV4_KV_SOURCE",    @"metal/dsv4_kv.metal"],
        @[@"DS4_METAL_DSV4_ROPE_SOURCE",  @"metal/dsv4_rope.metal"],
        @[@"DS4_METAL_DSV4_MISC_SOURCE",  @"metal/dsv4_misc.metal"],
        @[@"DS4_METAL_ARGSORT_SOURCE",    @"metal/argsort.metal"],
        @[@"DS4_METAL_CPY_SOURCE",        @"metal/cpy.metal"],
        @[@"DS4_METAL_CONCAT_SOURCE",     @"metal/concat.metal"],
        @[@"DS4_METAL_GET_ROWS_SOURCE",   @"metal/get_rows.metal"],
        @[@"DS4_METAL_SUM_ROWS_SOURCE",   @"metal/sum_rows.metal"],
        @[@"DS4_METAL_SOFTMAX_SOURCE",    @"metal/softmax.metal"],
        @[@"DS4_METAL_REPEAT_SOURCE",     @"metal/repeat.metal"],
        @[@"DS4_METAL_GLU_SOURCE",        @"metal/glu.metal"],
        @[@"DS4_METAL_NORM_SOURCE",       @"metal/norm.metal"],
        @[@"DS4_METAL_BIN_SOURCE",        @"metal/bin.metal"],
        @[@"DS4_METAL_SET_ROWS_SOURCE",   @"metal/set_rows.metal"],
    ];

    NSMutableString *source = [NSMutableString stringWithString:base];
    for (NSArray<NSString *> *spec in required_sources) {
        const char *override_path = getenv([spec[0] UTF8String]);
        NSMutableArray<NSString *> *paths = [NSMutableArray array];
        if (override_path && override_path[0]) {
            [paths addObject:[NSString stringWithUTF8String:override_path]];
        }
        [paths addObject:spec[1]];
        [paths addObject:[@"./" stringByAppendingString:spec[1]]];

        NSString *loaded = nil;
        NSString *loaded_path = nil;
        for (NSString *path in paths) {
            if (![fm fileExistsAtPath:path]) continue;

            NSError *error = nil;
            loaded = [NSString stringWithContentsOfFile:path
                                               encoding:NSUTF8StringEncoding
                                                  error:&error];
            if (!loaded) {
                fprintf(stderr, "ds4: failed to read Metal source %s: %s\n",
                        [path UTF8String], [[error localizedDescription] UTF8String]);
                return nil;
            }
            loaded_path = path;
            break;
        }

        if (!loaded) {
            fprintf(stderr,
                    "ds4: Metal source %s not found (set %s to override)\n",
                    [spec[1] UTF8String], [spec[0] UTF8String]);
            return nil;
        }
        [source appendFormat:@"\n// appended %@\n%@\n", loaded_path, loaded];
    }
    return source;
}

typedef struct {
    int32_t  ne00t;
    int32_t  ne00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne10;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
} ds4_gpu_get_rows_args;

typedef struct {
    int32_t  n_embd;
    int32_t  n_vocab;
    int32_t  n_tokens;
    uint64_t src_row_bytes;
    uint64_t dst_row_bytes;
    uint64_t token_stride;
} ds4_gpu_get_rows_q8_0_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    int32_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
} ds4_gpu_repeat_args;

typedef struct {
    int32_t  nk0;
    int32_t  ne01;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne11;
    int32_t  ne12;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
} ds4_gpu_set_rows_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    int32_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne10;
    int32_t  ne11;
    int32_t  ne12;
    int32_t  ne13;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    int32_t  dim;
} ds4_gpu_concat_args;

typedef struct {
    int64_t  nk0;
    int64_t  ne00;
    int64_t  ne01;
    int64_t  ne02;
    int64_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int64_t  ne0;
    int64_t  ne1;
    int64_t  ne2;
    int64_t  ne3;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
} ds4_gpu_cpy_args;

static ds4_gpu_cpy_args ds4_gpu_make_cpy_1d_args(
        uint32_t n,
        uint64_t src_elem,
        uint64_t dst_elem) {
    return (ds4_gpu_cpy_args) {
        .nk0 = (int64_t)n,
        .ne00 = (int64_t)n,
        .ne01 = 1,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = src_elem,
        .nb01 = (uint64_t)n * src_elem,
        .nb02 = (uint64_t)n * src_elem,
        .nb03 = (uint64_t)n * src_elem,
        .ne0 = (int64_t)n,
        .ne1 = 1,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = dst_elem,
        .nb1 = (uint64_t)n * dst_elem,
        .nb2 = (uint64_t)n * dst_elem,
        .nb3 = (uint64_t)n * dst_elem,
    };
}

static NSUInteger ds4_gpu_cpy_threads(uint32_t n, id<MTLComputePipelineState> pipeline) {
    NSUInteger nth = 32u;
    const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
    while (nth < (NSUInteger)n && nth < max_threads) nth *= 2u;
    if (nth > max_threads) nth = max_threads;
    if (nth > (NSUInteger)n) nth = (NSUInteger)n;
    return nth ? nth : 1u;
}

static float ds4_gpu_negative_infinity(void) {
    union { uint32_t u; float f; } v = { 0xff800000u };
    return v.f;
}

static float ds4_gpu_positive_infinity(void) {
    union { uint32_t u; float f; } v = { 0x7f800000u };
    return v.f;
}

static int ds4_gpu_encode_cpy_f32_f32_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n);

static int ds4_gpu_encode_cpy_f32_f32_3d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint32_t             planes,
        uint64_t             src_row_stride,
        uint64_t             src_plane_stride,
        uint64_t             dst_row_stride,
        uint64_t             dst_plane_stride);

static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint32_t             planes,
        uint64_t             src_col_stride,
        uint64_t             src_row_stride,
        uint64_t             src_plane_stride,
        uint64_t             dst_row_stride,
        uint64_t             dst_plane_stride);

static int ds4_gpu_encode_cpy_f32_f16_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n);

static int ds4_gpu_encode_cpy_f32_f16_2d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint64_t             src_row_stride,
        uint64_t             dst_row_stride);

static int ds4_gpu_encode_cpy_f16_f32_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n);

static int ds4_gpu_encode_fill_f32_rows(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        buf,
        NSUInteger           offset,
        uint32_t             width,
        uint32_t             rows,
        float                value);

static int ds4_gpu_encode_add_f32_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        a,
        NSUInteger           a_off,
        id<MTLBuffer>        b,
        NSUInteger           b_off,
        id<MTLBuffer>        out,
        NSUInteger           out_off,
        uint32_t             n);

typedef struct {
    int32_t  ne00;
    uint64_t nb01;
    int32_t  ne10;
    uint64_t nb11;
    int32_t  ne0;
    uint64_t nb1;
    int32_t  i00;
    int32_t  i10;
    float    alpha;
    float    limit;
} ds4_gpu_glu_args;

typedef struct {
    uint32_t n;
} ds4_gpu_add_flat_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    int32_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne10;
    int32_t  ne11;
    int32_t  ne12;
    int32_t  ne13;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    uint64_t offs;
    uint64_t o1[8];
} ds4_gpu_bin_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    int32_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    float    slope;
    float    scale;
    float    bias;
    float    val;
    float    min;
    float    max;
} ds4_gpu_unary_args;

static ds4_gpu_bin_args ds4_gpu_make_bin_rows_args(uint32_t n, uint32_t rows, uint32_t rhs_n) {
    const uint64_t row_bytes = (uint64_t)n * sizeof(float);
    const uint64_t rhs_row_bytes = (uint64_t)rhs_n * sizeof(float);
    return (ds4_gpu_bin_args) {
        .ne00 = (int32_t)n,
        .ne01 = (int32_t)rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = row_bytes,
        .nb02 = row_bytes,
        .nb03 = row_bytes,
        .ne10 = (int32_t)rhs_n,
        .ne11 = 1,
        .ne12 = 1,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = rhs_row_bytes,
        .nb12 = rhs_row_bytes,
        .nb13 = rhs_row_bytes,
        .ne0 = (int32_t)n,
        .ne1 = (int32_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = row_bytes,
        .nb2 = row_bytes,
        .nb3 = row_bytes,
        .offs = 0,
        .o1 = { 0 },
    };
}

static ds4_gpu_unary_args ds4_gpu_make_unary_rows_args(
        uint32_t n,
        uint32_t rows,
        int      c4,
        float    scale,
        float    bias) {
    const uint64_t row_bytes = (uint64_t)n * sizeof(float);
    const uint32_t n_kernel = c4 ? n / 4u : n;
    return (ds4_gpu_unary_args) {
        .ne00 = (int32_t)n_kernel,
        .ne01 = (int32_t)rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = row_bytes,
        .nb02 = row_bytes,
        .nb03 = row_bytes,
        .ne0 = (int32_t)n_kernel,
        .ne1 = (int32_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = row_bytes,
        .nb2 = row_bytes,
        .nb3 = row_bytes,
        .slope = 0.0f,
        .scale = scale,
        .bias = bias,
        .val = 0.0f,
        .min = 0.0f,
        .max = 0.0f,
    };
}

static ds4_gpu_bin_args ds4_gpu_make_bin_same_rows_args(uint32_t n, uint32_t rows) {
    const uint64_t row_bytes = (uint64_t)n * sizeof(float);
    return (ds4_gpu_bin_args) {
        .ne00 = (int32_t)n,
        .ne01 = (int32_t)rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = row_bytes,
        .nb02 = (uint64_t)rows * row_bytes,
        .nb03 = (uint64_t)rows * row_bytes,
        .ne10 = (int32_t)n,
        .ne11 = (int32_t)rows,
        .ne12 = 1,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = row_bytes,
        .nb12 = (uint64_t)rows * row_bytes,
        .nb13 = (uint64_t)rows * row_bytes,
        .ne0 = (int32_t)n,
        .ne1 = (int32_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = row_bytes,
        .nb2 = (uint64_t)rows * row_bytes,
        .nb3 = (uint64_t)rows * row_bytes,
        .offs = 0,
        .o1 = { 0 },
    };
}

static int ds4_gpu_encode_bin_f32_rows(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_bin_args   *args,
        id<MTLBuffer>               a,
        NSUInteger                  a_off,
        id<MTLBuffer>               b,
        NSUInteger                  b_off,
        id<MTLBuffer>               out,
        NSUInteger                  out_off);

static int ds4_gpu_encode_sum_rows_f32(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             width,
        uint32_t             rows);

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne10;
    int32_t  ne11;
    int32_t  ne12;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  nr0;
    int16_t  r2;
    int16_t  r3;
} ds4_gpu_q8_0_matvec_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne02;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne12;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ne0;
    int32_t  ne1;
    int16_t  r2;
    int16_t  r3;
} ds4_gpu_mul_mm_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne10;
    int32_t  ne11;
    int32_t  ne12;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ne0;
    int32_t  ne1;
    int16_t  r2;
    int16_t  r3;
} ds4_gpu_mul_mv_ext_args;

typedef ds4_gpu_q8_0_matvec_args ds4_gpu_f16_matvec_args;

static ds4_gpu_q8_0_matvec_args ds4_gpu_make_q8_0_mv_args(uint64_t in_dim, uint64_t out_dim) {
    const uint64_t row_bytes = (in_dim / 32u) * 34u;
    return (ds4_gpu_q8_0_matvec_args) {
        .ne00 = (int32_t)in_dim,
        .ne01 = (int32_t)out_dim,
        .ne02 = 1,
        .nb00 = 34,
        .nb01 = row_bytes,
        .nb02 = row_bytes * out_dim,
        .nb03 = row_bytes * out_dim,
        .ne10 = (int32_t)in_dim,
        .ne11 = 1,
        .ne12 = 1,
        .nb10 = sizeof(float),
        .nb11 = in_dim * sizeof(float),
        .nb12 = in_dim * sizeof(float),
        .nb13 = in_dim * sizeof(float),
        .ne0 = (int32_t)out_dim,
        .ne1 = 1,
        .nr0 = 2,
        .r2 = 1,
        .r3 = 1,
    };
}

static ds4_gpu_f16_matvec_args ds4_gpu_make_f16_mv_args(uint64_t in_dim, uint64_t out_dim) {
    const uint64_t row_bytes = in_dim * sizeof(uint16_t);
    return (ds4_gpu_f16_matvec_args) {
        .ne00 = (int32_t)in_dim,
        .ne01 = (int32_t)out_dim,
        .ne02 = 1,
        .nb00 = sizeof(uint16_t),
        .nb01 = row_bytes,
        .nb02 = row_bytes * out_dim,
        .nb03 = row_bytes * out_dim,
        .ne10 = (int32_t)in_dim,
        .ne11 = 1,
        .ne12 = 1,
        .nb10 = sizeof(float),
        .nb11 = in_dim * sizeof(float),
        .nb12 = in_dim * sizeof(float),
        .nb13 = in_dim * sizeof(float),
        .ne0 = (int32_t)out_dim,
        .ne1 = 1,
        .nr0 = 2,
        .r2 = 1,
        .r3 = 1,
    };
}

static ds4_gpu_q8_0_matvec_args ds4_gpu_make_f32_mv_args(
        uint64_t in_dim,
        uint64_t out_dim,
        uint64_t n_vec) {
    const uint64_t row_bytes = in_dim * sizeof(float);
    return (ds4_gpu_q8_0_matvec_args) {
        .ne00 = (int32_t)in_dim,
        .ne01 = (int32_t)out_dim,
        .ne02 = 1,
        .nb00 = sizeof(float),
        .nb01 = row_bytes,
        .nb02 = row_bytes * out_dim,
        .nb03 = row_bytes * out_dim,
        .ne10 = (int32_t)in_dim,
        .ne11 = (int32_t)n_vec,
        .ne12 = 1,
        .nb10 = sizeof(float),
        .nb11 = in_dim * sizeof(float),
        .nb12 = in_dim * n_vec * sizeof(float),
        .nb13 = in_dim * n_vec * sizeof(float),
        .ne0 = (int32_t)out_dim,
        .ne1 = (int32_t)n_vec,
        .nr0 = 2,
        .r2 = 1,
        .r3 = 1,
    };
}

typedef struct {
    const char *function_name;
    int16_t     nsg;
    int32_t     nr0;
    NSUInteger  smem;
} ds4_gpu_mv_dispatch;

static int ds4_gpu_tp_world_is_two(void);

static ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) {
    const uint64_t default_nsg = ds4_gpu_tp_world_is_two() ? 2u : 4u;
    const int16_t nsg =
        (int16_t)ds4_gpu_env_u64("DS4_METAL_Q8_MV_NSG", default_nsg, 1u, 8u);
    const uint64_t rows = ds4_gpu_env_u64("DS4_METAL_Q8_MV_ROWS", 2u, 2u, 4u);
    if (rows >= 4u) {
        return (ds4_gpu_mv_dispatch) {
            .function_name = "kernel_mul_mv_q8_0_f32_r4",
            .nsg = nsg,
            .nr0 = 4,
            .smem = 32u * 4u * sizeof(float),
        };
    }
    return (ds4_gpu_mv_dispatch) {
        .function_name = "kernel_mul_mv_q8_0_f32",
        .nsg = nsg,
        .nr0 = 2,
        .smem = 32u * 2u * sizeof(float),
    };
}

static ds4_gpu_mv_dispatch ds4_gpu_make_plain_mv_dispatch(
        uint64_t in_dim,
        int      f32_weights) {
    if (in_dim < 32) {
        return (ds4_gpu_mv_dispatch) {
            .function_name = f32_weights ? "kernel_mul_mv_f32_f32_short" : "kernel_mul_mv_f16_f32_short",
            .nsg = 1,
            .nr0 = 32,
            .smem = 0,
        };
    }

    const int16_t nsg = (int16_t)((in_dim + 127u) / 128u > 8u ? 8u : (in_dim + 127u) / 128u);
    const int use_4 = (in_dim % 4u) == 0;
    return (ds4_gpu_mv_dispatch) {
        .function_name = f32_weights
            ? (use_4 ? "kernel_mul_mv_f32_f32_4" : "kernel_mul_mv_f32_f32")
            : (use_4 ? "kernel_mul_mv_f16_f32_4" : "kernel_mul_mv_f16_f32"),
        .nsg = nsg,
        .nr0 = 2,
        .smem = 32u * 2u * sizeof(float),
    };
}

static ds4_gpu_mul_mm_args ds4_gpu_make_mm_args(
        uint64_t in_dim,
        uint64_t out_dim,
        uint64_t n_tok,
        uint64_t row_bytes) {
    return (ds4_gpu_mul_mm_args) {
        .ne00 = (int32_t)in_dim,
        .ne02 = 1,
        .nb01 = row_bytes,
        .nb02 = row_bytes * out_dim,
        .nb03 = row_bytes * out_dim,
        .ne12 = 1,
        .nb10 = sizeof(float),
        .nb11 = in_dim * sizeof(float),
        .nb12 = in_dim * n_tok * sizeof(float),
        .nb13 = in_dim * n_tok * sizeof(float),
        .ne0 = (int32_t)out_dim,
        .ne1 = (int32_t)n_tok,
        .r2 = 1,
        .r3 = 1,
    };
}

static ds4_gpu_mul_mv_ext_args ds4_gpu_make_mv_ext_args(
        uint64_t in_dim,
        uint64_t out_dim,
        uint64_t n_tok,
        uint64_t elem_bytes,
        uint64_t row_bytes) {
    return (ds4_gpu_mul_mv_ext_args) {
        .ne00 = (int32_t)in_dim,
        .ne01 = (int32_t)out_dim,
        .ne02 = 1,
        .nb00 = elem_bytes,
        .nb01 = row_bytes,
        .nb02 = row_bytes * out_dim,
        .nb03 = row_bytes * out_dim,
        .ne10 = (int32_t)in_dim,
        .ne11 = (int32_t)n_tok,
        .ne12 = 1,
        .nb10 = sizeof(float),
        .nb11 = in_dim * sizeof(float),
        .nb12 = in_dim * n_tok * sizeof(float),
        .nb13 = in_dim * n_tok * sizeof(float),
        .ne0 = (int32_t)out_dim,
        .ne1 = (int32_t)n_tok,
        .r2 = 1,
        .r3 = 1,
    };
}

static int16_t ds4_gpu_mv_ext_nxpsg(uint64_t in_dim, uint64_t n_tok) {
    if ((in_dim % 256u) == 0 && n_tok < 3) return 16;
    if ((in_dim % 128u) == 0) return 8;
    return 4;
}

static int16_t ds4_gpu_mv_ext_r1ptg(uint64_t n_tok) {
    switch (n_tok) {
    case 2: return 2;
    case 3:
    case 6: return 3;
    case 4:
    case 7:
    case 8: return 4;
    case 5: return 5;
    default: return n_tok > 8 ? 4 : 0;
    }
}

static const char *ds4_gpu_mv_ext_name(int q8, int16_t r1ptg) {
    if (q8) {
        switch (r1ptg) {
        case 2: return "kernel_mul_mv_ext_q8_0_f32_r1_2";
        case 3: return "kernel_mul_mv_ext_q8_0_f32_r1_3";
        case 4: return "kernel_mul_mv_ext_q8_0_f32_r1_4";
        case 5: return "kernel_mul_mv_ext_q8_0_f32_r1_5";
        default: return NULL;
        }
    }

    switch (r1ptg) {
    case 2: return "kernel_mul_mv_ext_f16_f32_r1_2";
    case 3: return "kernel_mul_mv_ext_f16_f32_r1_3";
    case 4: return "kernel_mul_mv_ext_f16_f32_r1_4";
    case 5: return "kernel_mul_mv_ext_f16_f32_r1_5";
    default: return NULL;
    }
}

static const char *ds4_gpu_mv_ext_f32_name(int16_t r1ptg) {
    switch (r1ptg) {
    case 2: return "kernel_mul_mv_ext_f32_f32_r1_2";
    case 3: return "kernel_mul_mv_ext_f32_f32_r1_3";
    case 4: return "kernel_mul_mv_ext_f32_f32_r1_4";
    case 5: return "kernel_mul_mv_ext_f32_f32_r1_5";
    default: return NULL;
    }
}

static const char *ds4_gpu_mv_ext_q8_pair_swiglu_name(int16_t r1ptg) {
    switch (r1ptg) {
    case 2: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2";
    case 3: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3";
    case 4: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4";
    case 5: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5";
    default: return NULL;
    }
}

typedef struct {
    int32_t  ne00;
    int32_t  ne00_t;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    float    eps;
    int32_t  nef1[3];
    int32_t  nef2[3];
    int32_t  nef3[3];
    uint64_t nbf1[3];
    uint64_t nbf2[3];
    uint64_t nbf3[3];
} ds4_gpu_rms_norm_args;

typedef struct {
    int32_t  q_n;
    int32_t  q_n4;
    int32_t  kv_n;
    int32_t  kv_n4;
    uint64_t q_row_stride;
    uint64_t kv_row_stride;
    float    eps;
} ds4_gpu_qkv_rms_norm_args;

static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_args(uint32_t n, uint32_t rows, float eps) {
    const uint64_t row_bytes = (uint64_t)n * sizeof(float);
    return (ds4_gpu_rms_norm_args) {
        .ne00 = (int32_t)n,
        .ne00_t = (int32_t)(n / 4u),
        .nb1 = row_bytes,
        .nb2 = row_bytes * rows,
        .nb3 = row_bytes * rows,
        .eps = eps,
        .nef1 = { (int32_t)rows, 1, 1 },
        .nef2 = { 1, 1, 1 },
        .nef3 = { 1, 1, 1 },
        .nbf1 = { row_bytes, row_bytes, row_bytes },
        .nbf2 = { row_bytes * rows, row_bytes, row_bytes },
        .nbf3 = { row_bytes * rows, row_bytes, row_bytes },
    };
}

static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_3d_args(
        uint32_t n0,
        uint32_t n1,
        uint32_t n2,
        float    eps) {
    const uint64_t row_bytes = (uint64_t)n0 * sizeof(float);
    const uint64_t plane_bytes = row_bytes * n1;
    return (ds4_gpu_rms_norm_args) {
        .ne00 = (int32_t)n0,
        .ne00_t = (int32_t)(n0 / 4u),
        .nb1 = row_bytes,
        .nb2 = plane_bytes,
        .nb3 = plane_bytes * n2,
        .eps = eps,
        .nef1 = { (int32_t)n1, 1, 1 },
        .nef2 = { (int32_t)n2, 1, 1 },
        .nef3 = { 1, 1, 1 },
        .nbf1 = { row_bytes, row_bytes, row_bytes },
        .nbf2 = { plane_bytes, row_bytes, row_bytes },
        .nbf3 = { plane_bytes * n2, row_bytes, row_bytes },
    };
}

static NSUInteger ds4_gpu_rms_norm_threads(uint32_t n) {
    NSUInteger ne00_t = n / 4u;
    NSUInteger nth = 32u;
    while (nth < ne00_t && nth < 1024u) nth *= 2u;
    if (nth > ne00_t) nth = ne00_t;
    return nth ? nth : 1u;
}

static NSUInteger ds4_gpu_rms_norm_pipeline_threads(
        uint32_t                  n,
        id<MTLComputePipelineState> pipeline) {
    NSUInteger ne00_t = n / 4u;
    NSUInteger max_threads = pipeline ? [pipeline maxTotalThreadsPerThreadgroup] : 1024u;
    NSUInteger nth = 32u;
    while (nth < ne00_t && nth < max_threads) nth *= 2u;
    if (nth > max_threads) nth = max_threads;
    if (nth > ne00_t) nth = ne00_t;
    return nth ? nth : 1u;
}

typedef struct {
    int32_t  n_hc;
    int32_t  sinkhorn_iters;
    int64_t  n_rows;
    int64_t  mix_hc;
    uint64_t nb01;
    uint64_t nb1;
    float    eps;
} ds4_gpu_hc_split_args;

typedef struct {
    int64_t n_embd;
    int64_t n_hc;
    int64_t n_tokens;
    uint64_t nb_x0;
    uint64_t nb_x1;
    uint64_t nb_x2;
    uint64_t nb_w0;
    uint64_t nb_w1;
    uint64_t nb0;
    uint64_t nb1;
} ds4_gpu_hc_weighted_sum_args;

typedef struct {
    int64_t n_embd;
    int64_t n_hc;
    int64_t n_tokens;
    uint64_t nb_x0;
    uint64_t nb_x1;
    uint64_t nb_x2;
    uint64_t nb_w0;
    uint64_t nb_w1;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb_norm1;
    float norm_eps;
} ds4_gpu_hc_weighted_sum_norm_args;

typedef struct {
    float post_scale;
    float eps;
} ds4_gpu_output_hc_weights4_args;

typedef struct {
    int64_t n_embd;
    int32_t n_hc;
    int32_t sinkhorn_iters;
    int64_t n_rows;
    int64_t mix_hc;
    uint64_t nb_mix1;
    uint64_t nb_split1;
    uint64_t nb_x0;
    uint64_t nb_x1;
    uint64_t nb_x2;
    uint64_t nb0;
    uint64_t nb1;
    float eps;
} ds4_gpu_hc_split_weighted_sum_args;

typedef struct {
    int64_t n_embd;
    int32_t n_hc;
    int32_t sinkhorn_iters;
    int64_t n_rows;
    int64_t mix_hc;
    uint64_t nb_mix1;
    uint64_t nb_split1;
    uint64_t nb_x0;
    uint64_t nb_x1;
    uint64_t nb_x2;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb_norm1;
    float eps;
    float norm_eps;
} ds4_gpu_hc_split_weighted_sum_norm_args;

typedef struct {
    int64_t n_embd;
    int64_t n_hc;
    int64_t n_tokens;
    uint64_t nb_block0;
    uint64_t nb_block1;
    uint64_t nb_add0;
    uint64_t nb_add1;
    uint64_t nb_res0;
    uint64_t nb_res1;
    uint64_t nb_res2;
    uint64_t nb_post0;
    uint64_t nb_post1;
    uint64_t nb_comb0;
    uint64_t nb_comb1;
    uint64_t nb_comb2;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    int32_t has_add;
} ds4_gpu_hc_expand_args;

typedef struct {
    int32_t  nei0;
    int32_t  nei1;
    uint64_t nbi1;
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    int32_t  ne10;
    int32_t  ne11;
    int32_t  ne12;
    int32_t  ne13;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    int32_t  ne0;
    int32_t  ne1;
    uint64_t nb1;
    int32_t  nr0;
    /* Tensor-parallel expert ownership; see ds4_metal_args_mul_mv_id in
     * metal/moe.metal.  Zero (from struct literals) means no split. */
    int32_t  tp_rank;
    int32_t  tp_world;
    int32_t  tp_addend;
    int32_t  tp_expert_base;
} ds4_gpu_mul_mv_id_args;

typedef struct {
    uint32_t n_total_expert;
    uint32_t n_expert;
} ds4_gpu_stream_expert_validate_args;

typedef struct {
    uint32_t active_mask;
    uint32_t accumulate;
} ds4_gpu_stream_expert_split_args;

typedef struct {
    int32_t  ne02;
    int32_t  ne10;
    int32_t  ne11;
    uint64_t nb11;
    uint64_t nb12;
    int32_t  ne21;
    int32_t  ne20;
    uint64_t nb21;
} ds4_gpu_mul_mm_id_map_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne02;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne11;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ne20;
    int32_t  ne21;
    int32_t  ne0;
    int32_t  ne1;
    int16_t  r2;
    int16_t  r3;
    int32_t  tp_rank;
    int32_t  tp_world;
    int32_t  tp_expert_base;
} ds4_gpu_mul_mm_id_args;

static int ds4_gpu_encode_mul_mv_id(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0);

static int ds4_gpu_encode_attn_out_low_q8_direct(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0);

static int ds4_gpu_encode_attn_out_low_q8_mpp(
        id<MTLCommandBuffer>           cb,
        id<MTLComputePipelineState>    pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>                  src0,
        NSUInteger                     src0_off,
        id<MTLBuffer>                  src1,
        NSUInteger                     src1_off,
        id<MTLBuffer>                  dst,
        NSUInteger                     dst_off);

static int ds4_gpu_encode_attn_out_low_q8_mpp(
        id<MTLCommandBuffer>           cb,
        id<MTLComputePipelineState>    pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>                  src0,
        NSUInteger                     src0_off,
        id<MTLBuffer>                  src1,
        NSUInteger                     src1_off,
        id<MTLBuffer>                  dst,
        NSUInteger                     dst_off);

static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args(
        uint32_t src0_cols,
        uint32_t src0_experts,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens);

static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args(
        uint32_t src0_cols,
        uint32_t src0_rows,
        uint32_t src0_experts,
        uint64_t src0_row_bytes,
        uint64_t src0_expert_bytes,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens);
static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size(
        uint32_t src0_cols,
        uint32_t src0_rows,
        uint32_t src0_experts,
        uint64_t src0_row_bytes,
        uint64_t src0_expert_bytes,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens,
        uint32_t src1_elem_size);

static int ds4_gpu_encode_mul_mm_id(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> map_pipeline,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_map_args *map_args,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off);

static int ds4_gpu_encode_mul_mm_id_map(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> map_pipeline,
        const ds4_gpu_mul_mm_id_map_args *map_args,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off);

static int ds4_gpu_encode_mul_mm_id_mapped(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off);
static int ds4_gpu_encode_mul_mm_id_mapped_tile(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes);
static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0_addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes,
        ds4_gpu_stream_expert_cache_entry * const *resources,
        uint32_t                    resource_count,
        uint32_t                    resource_kind,
        id<MTLBuffer>               overflow_resource);

typedef struct {
    int32_t  ne11;
    int32_t  ne_12_2;
    int32_t  ne_12_3;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    uint64_t nb21;
    uint64_t nb22;
    uint64_t nb23;
    int32_t  ne31;
    int32_t  ne32;
    int32_t  ne33;
    uint64_t nb31;
    uint64_t nb32;
    uint64_t nb33;
} ds4_gpu_flash_attn_pad_args;

typedef struct {
    uint32_t raw_cap;
    uint32_t raw_start;
    uint32_t n_raw;
    uint32_t n_comp;
    uint32_t pad_rows;
    uint32_t shared_pad;
} ds4_gpu_flash_kv_stage_f16_args;

typedef struct {
    int32_t  ne01;
    int32_t  ne30;
    int32_t  ne31;
    int32_t  ne32;
    int32_t  ne33;
    uint64_t nb31;
    uint64_t nb32;
    uint64_t nb33;
} ds4_gpu_flash_attn_blk_args;

typedef struct {
    int32_t  ne01;
    int32_t  ne02;
    int32_t  ne03;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne11;
    int32_t  ne_12_2;
    int32_t  ne_12_3;
    int32_t  ns10;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    int32_t  ns20;
    uint64_t nb21;
    uint64_t nb22;
    uint64_t nb23;
    int32_t  ne31;
    int32_t  ne32;
    int32_t  ne33;
    uint64_t nb31;
    uint64_t nb32;
    uint64_t nb33;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    float    scale;
    float    max_bias;
    float    m0;
    float    m1;
    int32_t  n_head_log2;
    float    logit_softcap;
} ds4_gpu_flash_attn_vec_args;

typedef struct {
    int32_t nrows;
} ds4_gpu_flash_attn_reduce_args;

typedef struct {
    int64_t ne00;
    int64_t ne01;
    int64_t ne02;
    int64_t ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    int32_t n_dims;
    int32_t mode;
    int32_t n_ctx_orig;
    int32_t inverse;
    float freq_base;
    float freq_scale;
    float ext_factor;
    float attn_factor;
    float beta_fast;
    float beta_slow;
    bool src2;
} ds4_gpu_rope_tail_batch_args;

typedef struct {
    uint64_t row_bytes;
    uint64_t token_bytes;
    int32_t head_dim;
    int32_t n_dims;
    int32_t n_ctx_orig;
    int32_t inverse;
    uint32_t pos0;
    uint32_t pos_step;
    float freq_base;
    float freq_scale;
    float ext_factor;
    float attn_factor;
    float beta_fast;
    float beta_slow;
} ds4_gpu_rope_affine_pair_args;

_Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64,
               "Metal affine RoPE argument ABI changed");

static ds4_gpu_rope_tail_batch_args ds4_gpu_make_rope_tail_args(
        uint32_t n_tok,
        uint32_t n_head,
        uint32_t head_dim,
        uint32_t n_rot,
        uint32_t n_ctx_orig,
        bool     inverse,
        float    freq_base,
        float    freq_scale,
        float    ext_factor,
        float    attn_factor,
        float    beta_fast,
        float    beta_slow) {
    const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
    const uint64_t tok_bytes = (uint64_t)n_head * row_bytes;
    return (ds4_gpu_rope_tail_batch_args) {
        .ne00 = head_dim,
        .ne01 = n_head,
        .ne02 = n_tok,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = row_bytes,
        .nb02 = tok_bytes,
        .nb03 = (uint64_t)n_tok * tok_bytes,
        .nb0 = sizeof(float),
        .nb1 = row_bytes,
        .nb2 = tok_bytes,
        .nb3 = (uint64_t)n_tok * tok_bytes,
        .n_dims = (int32_t)n_rot,
        .mode = 0,
        .n_ctx_orig = (int32_t)n_ctx_orig,
        .inverse = inverse ? 1 : 0,
        .freq_base = freq_base,
        .freq_scale = freq_scale,
        .ext_factor = ext_factor,
        .attn_factor = attn_factor,
        .beta_fast = beta_fast,
        .beta_slow = beta_slow,
        .src2 = false,
    };
}

static int ds4_gpu_encode_rope_tail_inplace(
        id<MTLCommandBuffer>                 cb,
        id<MTLBuffer>                        xbuf,
        NSUInteger                           xoff,
        const ds4_gpu_rope_tail_batch_args *args,
        uint32_t                             n_tok,
        uint32_t                             n_head,
        uint32_t                             head_dim,
        uint32_t                             pos0,
        uint32_t                             pos_step) {
    const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u;
    const bool lane_compatible =
        tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u;
    const bool force_affine_position =
        getenv("DS4_METAL_ENABLE_AFFINE_ROPE_PAIR") != NULL;
    const bool use_inplace_pair =
        g_rope_tail_inplace_pair_pipeline != nil &&
        args->mode == 0 && !args->src2 &&
        lane_compatible &&
        (ds4_gpu_device_name_contains("M3") ||
         (ds4_gpu_device_name_contains("M5") && n_tok == 1u) ||
         getenv("DS4_METAL_ENABLE_INPLACE_ROPE_PAIR") != NULL ||
         force_affine_position) &&
        getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL;
    const bool use_shared_coeff =
        use_inplace_pair && !force_affine_position &&
        g_rope_tail_inplace_pair_shared4_pipeline != nil &&
        /* The 256-thread grouped schedule helps long prefill, but reduces the
         * per-head parallelism that short batches and decode rely on. */
        tail_threads == 64u && n_head >= 4u && n_tok >= 32u &&
        getenv("DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF") == NULL;
    if (force_affine_position &&
        g_rope_tail_inplace_pair_affine_pipeline == nil &&
        getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL &&
        getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL) {
        fprintf(stderr,
                "ds4: forced affine-position RoPE pipeline is unavailable\n");
        return 0;
    }
    /* Keep long prefill on the proven shared4 kernel. Reconstructing affine
     * positions inside its coefficient cohort perturbs YaRN fast-math codegen;
     * the compact affine specialization is exact for the decode pair schedule. */
    const bool use_affine_position =
        use_inplace_pair && !use_shared_coeff &&
        g_rope_tail_inplace_pair_affine_pipeline != nil &&
        (n_tok == 1u || force_affine_position) &&
        (ds4_gpu_device_name_contains("M3") ||
         ds4_gpu_device_name_contains("M5") ||
         force_affine_position) &&
        getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL;

    int32_t pos_stack[256];
    int32_t *pos = NULL;
    id<MTLBuffer> posbuf = nil;
    const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(int32_t);
    if (!use_affine_position) {
        pos = pos_stack;
        if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) {
            pos = malloc((size_t)n_tok * sizeof(*pos));
            if (!pos) {
                fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n");
                return 0;
            }
        }
        for (uint32_t t = 0; t < n_tok; t++) {
            pos[t] = (int32_t)(pos0 + t * pos_step);
        }

        if (pos_bytes > 4096u) {
            /*
             * Metal inline setBytes data is meant for small constants. Long
             * prefill RoPE calls need thousands of positions; passing that much
             * inline can make the Apple driver abort the process.
             */
            posbuf = ds4_gpu_new_transient_buffer(
                pos_bytes, "ds4_rope_positions");
            if (!posbuf) {
                if (pos != pos_stack) free(pos);
                return 0;
            }
            memcpy([posbuf contents], pos, pos_bytes);
        }
    }

    ds4_gpu_rope_affine_pair_args affine_args;
    if (use_affine_position) {
        const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
        affine_args = (ds4_gpu_rope_affine_pair_args) {
            .row_bytes = row_bytes,
            .token_bytes = (uint64_t)n_head * row_bytes,
            .head_dim = (int32_t)head_dim,
            .n_dims = args->n_dims,
            .n_ctx_orig = args->n_ctx_orig,
            .inverse = args->inverse,
            .pos0 = pos0,
            .pos_step = pos_step,
            .freq_base = args->freq_base,
            .freq_scale = args->freq_scale,
            .ext_factor = args->ext_factor,
            .attn_factor = args->attn_factor,
            .beta_fast = args->beta_fast,
            .beta_slow = args->beta_slow,
        };
    }
    const NSUInteger reference_nth =
        (NSUInteger)(head_dim < 256u ? head_dim : 256u);
    const NSUInteger pair_nth =
        (NSUInteger)(tail_threads < 256u ? tail_threads : 256u);
    const NSUInteger nth = use_shared_coeff ? 256u :
        (use_inplace_pair ? pair_nth : reference_nth);
    const NSUInteger head_groups = use_shared_coeff ?
        (NSUInteger)((n_head + 3u) / 4u) : (NSUInteger)n_head;
    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    id<MTLComputePipelineState> pipeline = use_affine_position ?
        g_rope_tail_inplace_pair_affine_pipeline :
        (use_shared_coeff ?
            g_rope_tail_inplace_pair_shared4_pipeline :
            (use_inplace_pair ?
                g_rope_tail_inplace_pair_pipeline : g_rope_tail_batch_pipeline));
    [enc setComputePipelineState:pipeline];
    if (use_affine_position) {
        [enc setBytes:&affine_args length:sizeof(affine_args) atIndex:0];
    } else {
        [enc setBytes:args length:sizeof(*args) atIndex:0];
    }
    [enc setBuffer:xbuf offset:xoff atIndex:1];
    if (!use_affine_position) {
        if (posbuf) {
            [enc setBuffer:posbuf offset:0 atIndex:2];
        } else {
            [enc setBytes:pos length:pos_bytes atIndex:2];
        }
        [enc setBuffer:xbuf offset:xoff atIndex:3];
    }
    [enc setBuffer:xbuf offset:xoff atIndex:4];
    [enc dispatchThreadgroups:MTLSizeMake(head_groups, n_tok, 1)
         threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    if (pos && pos != pos_stack) free(pos);
    return 1;
}

typedef struct {
    int64_t ne00;
    int64_t ne01;
    int64_t ne02;
    int64_t ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    int32_t n_rot;
} ds4_gpu_dsv4_fp8_kv_quantize_args;

typedef struct {
    int32_t head_dim;
    int32_t n_rot;
    int32_t raw_row;
} ds4_gpu_dsv4_kv_fp8_store_args;

typedef struct {
    uint32_t n_rows;
    uint32_t head_dim;
    uint64_t row_stride;
} ds4_gpu_dsv4_indexer_qat_args;

typedef struct {
    uint32_t width;
} ds4_gpu_dsv4_ratio4_shift_args;

typedef struct {
    uint32_t head_dim;
    uint32_t n_comp;
    uint32_t replay;
    uint32_t n_threads;
} ds4_gpu_dsv4_compressor_pack_ratio4_args;

typedef struct {
    int64_t  n_rows;
    uint32_t head_dim;
    uint32_t n_comp;
    uint32_t replay;
    uint32_t pad;
} ds4_gpu_dsv4_softmax_pool_ratio4_direct_args;

typedef struct {
    uint32_t width;
    uint32_t ratio;
    uint32_t pos;
    uint32_t ape_type;
} ds4_gpu_dsv4_compressor_store_one_args;

typedef struct {
    int64_t  ne00;
    int64_t  ne01;
    int64_t  ne02;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb10;
    uint64_t nb11;
    uint64_t nb12;
    int64_t  ne0;
    int64_t  ne1;
    uint64_t nb0;
    uint64_t nb1;
} ds4_gpu_dsv4_softmax_pool_args;

typedef struct {
    uint32_t width;
    uint32_t ratio;
    uint32_t pos0;
    uint32_t n_tokens;
} ds4_gpu_dsv4_compressor_score_ape_args;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    int32_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    int32_t  top_k;
} ds4_gpu_kargs_argsort;

typedef struct {
    int64_t  ne00;
    int64_t  ne01;
    int64_t  ne02;
    int64_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne0;
    int32_t  ne1;
    int32_t  ne2;
    int32_t  ne3;
    int32_t  top_k;
    int32_t  len;
} ds4_gpu_kargs_argsort_merge;

typedef struct {
    int64_t  ne00;
    int64_t  ne01;
    int64_t  ne02;
    int64_t  ne03;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int64_t  ne0;
    int64_t  ne1;
    int64_t  ne2;
    int64_t  ne3;
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
} ds4_gpu_kargs_sum_rows;

typedef struct {
    int32_t  ne00;
    int32_t  ne01;
    int32_t  ne02;
    uint64_t nb01;
    uint64_t nb02;
    uint64_t nb03;
    int32_t  ne11;
    int32_t  ne12;
    int32_t  ne13;
    uint64_t nb11;
    uint64_t nb12;
    uint64_t nb13;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    float    scale;
    float    max_bias;
    float    m0;
    float    m1;
    int32_t  n_head_log2;
} ds4_gpu_softmax_args;

typedef struct {
    int64_t  ne00;
    int64_t  ne01;
    uint64_t nb00;
    uint64_t nb01;
    int64_t  ne0;
    int64_t  ne1;
    uint64_t nb0;
    uint64_t nb1;
} ds4_gpu_dsv4_topk_mask_args;

typedef struct {
    int64_t  ne00;
    int64_t  ne01;
    int64_t  ne02;
    uint64_t nb00;
    uint64_t nb01;
    uint64_t nb02;
    int64_t  ne10;
    int64_t  ne11;
    uint64_t nb10;
    uint64_t nb11;
    int64_t  ne0;
    int64_t  ne1;
    uint64_t nb0;
    uint64_t nb1;
    float    scale;
} ds4_gpu_dsv4_indexer_weighted_sum_args;

typedef struct {
    uint32_t has_bias;
    uint32_t hash_mode;
    uint32_t use_token_buffer;
    uint32_t token;
    uint32_t hash_rows;
} ds4_gpu_dsv4_router_select_one_args;

typedef struct {
    uint32_t n_expert;
    uint32_t n_expert_used;
    float    expert_weight_scale;
    uint32_t pad0;
} ds4_gpu_glm_router_select_one_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t kv_raw_dim;
    uint32_t kv_lora_dim;
    float    eps;
} ds4_gpu_glm_kv_lora_rms_norm_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t n_head;
    uint32_t row_bytes;
    uint32_t weight_type;
    uint32_t pad1;
    uint32_t pad2;
} ds4_gpu_glm_k_b_project_args;

typedef struct {
    uint32_t pos0;
    uint32_t n_tokens;
    uint32_t cache_cap;
    uint32_t kv_raw_dim;
    uint32_t kv_lora_dim;
    uint32_t qk_rope;
    uint32_t cache_f16;
    uint32_t pad1;
} ds4_gpu_glm_store_compact_kv_args;

typedef struct {
    uint32_t pos0;
    uint32_t n_tokens;
    uint32_t cache_cap;
    uint32_t q_n;
    uint32_t q_n4;
    uint32_t kv_raw_dim;
    uint32_t kv_lora_dim;
    uint32_t kv_lora_n4;
    uint32_t qk_rope;
    uint32_t cache_f16;
    float    eps;
    uint32_t pad0;
} ds4_gpu_glm_qkv_norm_store_compact_kv_args;

typedef struct {
    uint32_t pos0;
    uint32_t n_tokens;
    uint32_t cache_cap;
    uint32_t head_dim;
    uint32_t rot_dim;
    uint32_t n_ctx_orig;
    uint32_t cache_f16;
    uint32_t pad0;
    float    eps;
    float    freq_base;
    float    freq_scale;
    float    ext_factor;
    float    attn_factor;
    float    beta_fast;
    float    beta_slow;
    float    pad1;
} ds4_gpu_glm_store_indexer_k_args;

typedef struct {
    uint32_t pos0;
    uint32_t n_tokens;
    uint32_t cache_cap;
    uint32_t n_head;
    uint32_t kv_raw_dim;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t qk_rope;
    uint32_t value_dim;
    uint32_t n_ctx_orig;
    uint32_t cache_f16;
    uint32_t pad0;
    float    freq_base;
    float    freq_scale;
    float    ext_factor;
    float    attn_factor;
    float    beta_fast;
    float    beta_slow;
} ds4_gpu_glm_build_kv_cache_args;

typedef struct {
    uint32_t pos0;
    uint32_t n_tokens;
    uint32_t cache_len;
    uint32_t cache_cap;
    uint32_t n_head;
    uint32_t qk_dim;
    uint32_t value_dim;
    uint32_t pad0;
    uint32_t cache_f16;
    uint32_t pad1;
    uint32_t pad2;
    float    scale;
} ds4_gpu_glm_attention_full_args;

typedef struct {
    uint32_t n_selected;
} ds4_gpu_glm_fill_selected_range_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t pos0;
    uint32_t n_selected;
    uint32_t pad_row;
} ds4_gpu_glm_fill_selected_range_batch_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t n_head;
    uint32_t head_dim;
    uint32_t rot_dim;
    uint32_t rot_offset;
    uint32_t pos0;
    uint32_t n_ctx_orig;
    float    freq_base;
    float    freq_scale;
    float    ext_factor;
    float    attn_factor;
    float    beta_fast;
    float    beta_slow;
} ds4_gpu_glm_rope_tail_args;

typedef struct {
    uint32_t n_rows;
    uint32_t n_head;
    uint32_t head_dim;
    uint32_t cache_f16;
    float    scale;
} ds4_gpu_glm_indexer_score_one_args;

typedef struct {
    uint32_t n_rows;
    uint32_t n_tokens;
    uint32_t n_head;
    uint32_t head_dim;
    uint32_t pos0;
    uint32_t cache_f16;
    uint64_t q_token_stride;
    uint64_t q_head_stride;
    uint64_t weights_token_stride;
    uint64_t score_token_stride;
    float    scale;
} ds4_gpu_glm_indexer_scores_batch_args;

typedef struct {
    uint32_t n_head;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t qk_dim;
    uint32_t row_bytes;
    uint32_t weight_type;
    uint32_t pad1;
    uint32_t pad2;
} ds4_gpu_glm_qk_lowrank_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t n_head;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t qk_dim;
    uint32_t row_bytes;
    uint32_t weight_type;
    uint32_t head_base;
} ds4_gpu_glm_qk_lowrank_batch_args;

typedef struct {
    uint32_t n_selected;
    uint32_t cache_cap;
    uint32_t cache_f16;
    uint32_t n_head;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t qk_rope;
    uint32_t value_dim;
    uint32_t n_ctx_orig;
    uint32_t value_row_bytes;
    float    scale;
    float    freq_base;
    float    freq_scale;
    float    ext_factor;
    float    attn_factor;
    float    beta_fast;
    float    beta_slow;
    uint32_t value_type;
} ds4_gpu_glm_attention_indexed_decode_args;

typedef struct {
    uint32_t n_selected;
    uint32_t cache_cap;
    uint32_t cache_f16;
    uint32_t n_head;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t qk_rope;
    uint32_t value_dim;
    uint32_t n_ctx_orig;
    uint32_t value_row_bytes;
    uint32_t block_rows;
    uint32_t n_blocks;
    float    scale;
    float    freq_base;
    float    freq_scale;
    float    ext_factor;
    float    attn_factor;
    float    beta_fast;
    float    beta_slow;
    uint32_t value_type;
} ds4_gpu_glm_attention_indexed_decode_split_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t n_selected;
    uint32_t cache_cap;
    uint32_t cache_f16;
    uint32_t n_head;
    uint32_t kv_lora_dim;
    uint32_t qk_nope;
    uint32_t qk_rope;
    uint32_t value_dim;
    uint32_t n_ctx_orig;
    uint32_t value_row_bytes;
    uint32_t value_type;
    uint32_t pos0;
    float    scale;
    float    freq_base;
    float    freq_scale;
    float    ext_factor;
    float    attn_factor;
    float    beta_fast;
    float    beta_slow;
    uint32_t head_base;
} ds4_gpu_glm_attention_indexed_batch_args;

typedef struct {
    uint32_t in_dim;
    uint32_t mid_dim;
    uint32_t out_dim;
    uint32_t n_total_expert;
    uint32_t n_expert_used;
    uint32_t n_tokens;
    uint32_t mid_token_stride;
    uint32_t down_type;
    int32_t  tp_rank;
    int32_t  tp_world;
    int32_t  tp_expert_base;
    uint64_t gate_expert_bytes;
    uint64_t gate_row_bytes;
    uint64_t up_expert_bytes;
    uint64_t up_row_bytes;
    uint64_t down_expert_bytes;
    uint64_t down_row_bytes;
} ds4_gpu_glm_routed_moe_args;

typedef struct {
    uint32_t n_tokens;
    uint32_t n_head;
    uint32_t n_raw;
    uint32_t raw_cap;
    uint32_t raw_start;
    uint32_t n_comp;
    uint32_t top_k;
    uint32_t pos0;
    uint32_t window;
    uint32_t ratio;
    uint32_t comp_kv_f16;
    uint32_t pad0;
    uint64_t q_token_stride;
    uint64_t q_head_stride;
    uint64_t raw_row_stride;
    uint64_t comp_row_stride;
    uint64_t topk_token_stride;
    uint64_t dst_token_stride;
    uint64_t dst_head_stride;
    float    scale;
} ds4_gpu_dsv4_indexed_attention_args;

typedef struct {
    uint32_t n_comp;
    uint32_t n_tokens;
    uint32_t n_head;
    uint32_t head_dim;
    uint32_t pos0;
    uint32_t ratio;
    uint64_t q_token_stride;
    uint64_t q_head_stride;
    uint64_t weights_token_stride;
    uint64_t index_row_stride;
    uint64_t score_token_stride;
    float    scale;
} ds4_gpu_dsv4_indexer_scores_fused_args;

typedef struct {
    uint32_t width;
    uint32_t rows;
    uint64_t gate_row_stride;
    uint64_t up_row_stride;
    uint64_t mid_row_stride;
    uint64_t weight_stride;
    uint32_t write_clamped;
    float    clamp_value;
} ds4_gpu_dsv4_moe_swiglu_weight_args;

typedef struct {
    uint32_t expert_base;
    uint32_t expert_count;
    uint32_t accumulate;
    uint32_t pad0;
} ds4_gpu_moe_expert_group_args;

typedef struct {
    uint64_t expert_bytes;
    uint32_t group_size;
    uint32_t n_slots;
} ds4_gpu_q4_gather_slots6_args;

typedef struct {
    uint32_t width;
    uint32_t tokens;
    uint64_t src_token_stride;
    uint64_t dst_token_stride;
} ds4_gpu_dsv4_moe_sum6_args;

/* Compile the single in-repo Metal source and create the pipelines that every
 * session uses. Shape-dependent kernels with function constants are built
 * lazily by the small ds4_gpu_get_* caches, so startup stays predictable
 * while long-context prefill and decode can still pick specialized variants. */
int ds4_gpu_init(void) {
    if (g_initialized) return 1;

    @autoreleasepool {
        g_device = MTLCreateSystemDefaultDevice();
        if (!g_device) {
            fprintf(stderr, "ds4: Metal device not available\n");
            return 0;
        }
        ds4_gpu_print_device_summary();
        ds4_gpu_detect_metal4_features();

        g_queue = [g_device newCommandQueue];
        if (!g_queue) {
            fprintf(stderr, "ds4: failed to create Metal command queue\n");
            g_device = nil;
            return 0;
        }
        g_model_buffer_cache = [NSMutableDictionary dictionary];
        g_model_buffer_cache_bytes = 0;
        g_model_buffer_cache_evictions = 0;
        g_model_buffer_cache_over_limit = 0;
        g_q4_expert_table_cache = [NSMutableDictionary dictionary];
        g_q4_expert_layer_residency_cache = [NSMutableDictionary dictionary];
        g_pipeline_cache = [NSMutableDictionary dictionary];
        g_transient_buffers = [NSMutableArray array];
        g_pending_cbs = [NSMutableArray array];
        if (!g_model_buffer_cache || !g_q4_expert_table_cache ||
            !g_q4_expert_layer_residency_cache ||
            !g_pipeline_cache || !g_transient_buffers || !g_pending_cbs) {
            fprintf(stderr, "ds4: Metal bookkeeping allocation failed\n");
            g_pending_cbs = nil;
            g_transient_buffers = nil;
            g_pipeline_cache = nil;
            g_q4_expert_layer_residency_cache = nil;
            g_q4_expert_table_cache = nil;
            g_model_buffer_cache = nil;
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        NSError *error = nil;
        NSString *source = ds4_gpu_full_source();
        if (!source) {
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        MTLCompileOptions *options = [MTLCompileOptions new];
        NSMutableDictionary *macros = [NSMutableDictionary new];
        if (g_metal4_tensor_api_enabled) {
            macros[@"DS4_METAL_HAS_TENSOR"] = @"1";
            fprintf(stderr, "ds4: Metal 4 tensor API enabled for Tensor kernels\n");
        }

        const int drift_hc_stable        = ds4_gpu_env_bool("DS4_METAL_HC_STABLE")          != 0; // default ON
        const int drift_norm_unify       = ds4_gpu_env_bool("DS4_METAL_NORM_RSQRT_DISABLE") != 0; // default ON
        const int drift_kv_raw_f32       = ds4_gpu_env_bool("DS4_METAL_KV_RAW_F32")         >  0; // default OFF
        const int drift_rope_exp2_log2   = ds4_gpu_env_bool("DS4_METAL_ROPE_EXP2_LOG2")     >  0; // default OFF
        const int drift_math_safe        = ds4_gpu_env_bool("DS4_METAL_MATH_SAFE")          >  0; // default OFF

        if (drift_math_safe) {
            // MTLCompileOptions.fastMathEnabled defaults to YES and Apple's
            // headers explicitly say this "may violate the IEEE 754 standard".
            // Different fast-math optimizations get applied across the
            // matmul2d cooperative-tensor path and the legacy
            // simdgroup_multiply_accumulate path on M5, amplifying the
            // mismatch. MTLMathModeSafe pins the entire library to strict
            // IEEE-754 semantics. Diagnostic-only: useful to localize drift
            // sources but not to ship as a default.
            if (@available(macOS 15.0, *)) {
                options.mathMode = MTLMathModeSafe;
                fprintf(stderr, "ds4: Metal shader library math mode = safe (strict IEEE-754) by DS4_METAL_MATH_SAFE\n");
            } else {
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
                options.fastMathEnabled = NO;
#pragma clang diagnostic pop
                fprintf(stderr, "ds4: Metal shader library fast-math disabled by DS4_METAL_MATH_SAFE (pre-macOS 15)\n");
            }
        }

        if (drift_hc_stable)      macros[@"DS4_METAL_HC_STABLE"]          = @"1";
        if (drift_norm_unify)     macros[@"DS4_METAL_NORM_RSQRT_DISABLE"] = @"1";
        if (drift_kv_raw_f32)     macros[@"DS4_METAL_KV_RAW_F32"]         = @"1";
        if (drift_rope_exp2_log2) macros[@"DS4_METAL_ROPE_EXP2_LOG2"]     = @"1";
        fprintf(stderr,
                "ds4: drift-patch flags hc_stable=%s norm_unify=%s kv_raw_f32=%s rope_exp2_log2=%s math_safe=%s tensor_matmul=%s\n",
                drift_hc_stable      ? "on"  : "off",
                drift_norm_unify     ? "on"  : "off",
                drift_kv_raw_f32     ? "on"  : "off",
                drift_rope_exp2_log2 ? "on"  : "off",
                drift_math_safe      ? "on"  : "off",
                g_metal4_tensor_api_enabled ? "on" : "off");
        options.preprocessorMacros = macros;
        id<MTLLibrary> library = [g_device newLibraryWithSource:source options:options error:&error];
        if (!library) {
            fprintf(stderr, "ds4: Metal shader compilation failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_library = library;

        id<MTLFunction> fn = [library newFunctionWithName:@"kernel_get_rows_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_get_rows_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_get_rows_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_get_rows_f16"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_f16 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_get_rows_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_get_rows_f16_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_f16 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_get_rows_i32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_i32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_get_rows_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_get_rows_i32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_i32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_get_rows_q8_0_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_get_rows_q8_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_get_rows_q8_0_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_get_rows_q4_0_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_get_rows_q4_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_get_rows_q4_0_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_get_rows_q4_K_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_get_rows_q4_K_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_get_rows_q4_K_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_repeat_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_repeat_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_repeat_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_repeat_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_repeat_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_set_rows_f32_i32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_set_rows_f32_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_set_rows_f32_i32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_concat"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_concat function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_concat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_concat_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_concat pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_f32_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_f32_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_f32_f16"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_f32_f16_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_contig_f32_f16_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_contig_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_contig_f32_f16_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_f16_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_f16_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_f16_f16"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_f16_f16_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f32_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_contig_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_contig_f16_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f16_bits_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_cpy_contig_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_cpy_contig_f16_f16_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_flash_kv_stage_f16"];
        if (fn) {
            g_flash_kv_stage_f16_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_flash_kv_stage_f16_pipeline) {
                fprintf(stderr,
                        "ds4: optional Metal gathered KV staging pipeline unavailable: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        } else {
            fprintf(stderr,
                    "ds4: optional Metal gathered KV staging kernel unavailable\n");
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_fp8_kv_quantize_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_fp8_kv_quantize_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_fp8_kv_quantize_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_indexer_hadamard_fp4_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_indexer_qat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_indexer_qat_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_kv_fp8_store_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_kv_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_kv_fp8_store_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_ratio4_shift_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_ratio4_shift_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_ratio4_shift_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_swiglu_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_swiglu_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_swiglu_flat_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        g_swiglu_flat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_swiglu_flat_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum6_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_moe_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum8_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_moe_sum8_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_sum8_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *bin_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t bin_op = 0;
        int16_t bin_f = 1;
        bool bin_rb = false;
        bool bin_cb = false;
        [bin_constants setConstantValue:&bin_op type:MTLDataTypeShort atIndex:1300];
        [bin_constants setConstantValue:&bin_f  type:MTLDataTypeShort atIndex:1301];
        [bin_constants setConstantValue:&bin_rb type:MTLDataTypeBool  atIndex:1302];
        [bin_constants setConstantValue:&bin_cb type:MTLDataTypeBool  atIndex:1303];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
                           constantValues:bin_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_add_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_add_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_add2_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_add2_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        g_add2_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_add2_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_add2_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_add3_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_add3_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        g_add3_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_add3_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_add3_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *bin_mul_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t bin_mul_plain_op = 2;
        int16_t bin_mul_plain_f = 1;
        bool bin_mul_plain_rb = false;
        bool bin_mul_plain_cb = false;
        [bin_mul_constants setConstantValue:&bin_mul_plain_op type:MTLDataTypeShort atIndex:1300];
        [bin_mul_constants setConstantValue:&bin_mul_plain_f  type:MTLDataTypeShort atIndex:1301];
        [bin_mul_constants setConstantValue:&bin_mul_plain_rb type:MTLDataTypeBool  atIndex:1302];
        [bin_mul_constants setConstantValue:&bin_mul_plain_cb type:MTLDataTypeBool  atIndex:1303];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
                           constantValues:bin_mul_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_mul_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_mul_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *bin_mul_scalar_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t bin_mul_op = 2;
        int16_t bin_mul_f = 1;
        bool bin_mul_rb = false;
        bool bin_mul_cb = true;
        [bin_mul_scalar_constants setConstantValue:&bin_mul_op type:MTLDataTypeShort atIndex:1300];
        [bin_mul_scalar_constants setConstantValue:&bin_mul_f  type:MTLDataTypeShort atIndex:1301];
        [bin_mul_scalar_constants setConstantValue:&bin_mul_rb type:MTLDataTypeBool  atIndex:1302];
        [bin_mul_scalar_constants setConstantValue:&bin_mul_cb type:MTLDataTypeBool  atIndex:1303];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
                           constantValues:bin_mul_scalar_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_bin_mul_scalar_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_bin_mul_scalar_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *bin_div_row_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t bin_div_op = 3;
        int16_t bin_div_f = 1;
        bool bin_div_rb = false;
        bool bin_div_cb = true;
        [bin_div_row_constants setConstantValue:&bin_div_op type:MTLDataTypeShort atIndex:1300];
        [bin_div_row_constants setConstantValue:&bin_div_f  type:MTLDataTypeShort atIndex:1301];
        [bin_div_row_constants setConstantValue:&bin_div_rb type:MTLDataTypeBool  atIndex:1302];
        [bin_div_row_constants setConstantValue:&bin_div_cb type:MTLDataTypeBool  atIndex:1303];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
                           constantValues:bin_div_row_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_bin_div_row_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_bin_div_row_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_rms_norm_mul_f32_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_rms_norm_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_rms_norm_f32_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_rms_norm_plain_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_rms_norm_plain_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_add_rms_norm_mul_f32_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_add_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_add_rms_norm_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_qkv_rms_norm_f32_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_qkv_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_qkv_rms_norm_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *moe_mv_id_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t moe_mv_id_nsg = 2;
        [moe_mv_id_constants setConstantValue:&moe_mv_id_nsg type:MTLDataTypeShort atIndex:600];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_iq2_xxs_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_iq2_xxs_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_iq2_xxs_pair_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_q2_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_q2_k_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_q2_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_iq2_xxs_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q2_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_slots6_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_slots6_q2_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_addr_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_addr_iq2_xxs_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_addr_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_addr_q2_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline =
            [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_masked_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline =
            [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_stream_expert_cache_validate"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_stream_expert_cache_validate_pipeline =
            [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_stream_expert_cache_validate_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_q4_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_q4_k_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_q4_k_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_q4_k_pair_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_id_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_id_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group6_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group8_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group8_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_id_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group24_q4_k_id_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group24_q4_k_id_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_group24_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_group24_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_slots6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_slots6_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_q4_gather_slots6"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_q4_gather_slots6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_q4_gather_slots6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_table_q4_pair_gate_encoder = [fn newArgumentEncoderWithBufferIndex:2];
        g_moe_table_q4_pair_up_encoder = [fn newArgumentEncoderWithBufferIndex:3];
        if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder) {
            fprintf(stderr, "ds4: Metal Q4 expert-table pair argument encoder creation failed\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_table_q4_sum_down_encoder = [fn newArgumentEncoderWithBufferIndex:1];
        if (!g_moe_table_q4_sum_down_encoder) {
            fprintf(stderr, "ds4: Metal Q4 expert-table down argument encoder creation failed\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_moe_mul_mv_table_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_moe_mul_mv_table_q4_k_sum6_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_pair_swiglu_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (fn) {
            g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline) {
                fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 pipeline unavailable: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        } else {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 function unavailable: %s\n",
                    [[error localizedDescription] UTF8String]);
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_sum6_f32"
                           constantValues:moe_mv_id_constants
                                    error:&error];
        if (fn) {
            g_moe_mul_mv_addr_q4_k_sum6_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_moe_mul_mv_addr_q4_k_sum6_pipeline) {
                fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 pipeline unavailable: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        } else {
            fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 function unavailable: %s\n",
                    [[error localizedDescription] UTF8String]);
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_rope_tail_batch_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_rope_tail_batch_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_rope_tail_inplace_pair_pipeline =
            [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_rope_tail_inplace_pair_pipeline) {
            fprintf(stderr,
                    "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_shared4"];
        if (!fn) {
            fprintf(stderr,
                    "ds4: optional Metal shared-head RoPE kernel unavailable; using per-head path\n");
        } else {
            g_rope_tail_inplace_pair_shared4_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_rope_tail_inplace_pair_shared4_pipeline) {
                fprintf(stderr,
                        "ds4: optional Metal shared-head RoPE pipeline unavailable; using per-head path: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_affine"];
        if (fn) {
            g_rope_tail_inplace_pair_affine_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_rope_tail_inplace_pair_affine_pipeline) {
                fprintf(stderr,
                        "ds4: optional Metal affine-position RoPE pair pipeline unavailable: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        } else {
            fprintf(stderr,
                    "ds4: optional Metal affine-position RoPE pair kernel unavailable\n");
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_softmax_pool"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_softmax_pool_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_softmax_pool_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_soft_max_f32"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_soft_max_f32 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_soft_max_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_soft_max_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_soft_max_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_soft_max_f32_4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_soft_max_f32_4_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_soft_max_f32_4_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_argsort_f32_i32_desc"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_argsort_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_argsort_f32_i32_desc_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_argsort_merge_f32_i32_desc"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_argsort_merge_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_argsort_merge_f32_i32_desc_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *sum_rows_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t sum_rows_op = 10;
        [sum_rows_constants setConstantValue:&sum_rows_op type:MTLDataTypeShort atIndex:1400];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_sum_rows_f32_f32"
                           constantValues:sum_rows_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_sum_rows_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_sum_rows_f32_f32_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_topk_mask_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_topk_mask_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask_scatter"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_topk_mask_scatter_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_topk_mask_scatter_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_indexer_weighted_sum"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_dsv4_indexer_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_dsv4_indexer_weighted_sum_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_sinkhorn"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_hc_split_sinkhorn_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_hc_split_sinkhorn_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_hc_split_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_hc_split_weighted_sum_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum_norm4"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_hc_split_weighted_sum_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_hc_split_weighted_sum_norm_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_hc_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_hc_weighted_sum_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum_norm4"];
        if (fn) {
            g_hc_weighted_sum_norm_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_hc_weighted_sum_norm_pipeline) {
                fprintf(stderr,
                        "ds4: optional Metal output HC sum/RMSNorm pipeline unavailable: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        } else {
            fprintf(stderr,
                    "ds4: optional Metal output HC sum/RMSNorm kernel unavailable\n");
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"];
        if (fn) {
            g_output_hc_weights4_pipeline =
                [g_device newComputePipelineStateWithFunction:fn error:&error];
            if (!g_output_hc_weights4_pipeline) {
                fprintf(stderr,
                        "ds4: optional Metal output HC weights4 pipeline unavailable: %s\n",
                        [[error localizedDescription] UTF8String]);
            }
        } else {
            fprintf(stderr,
                    "ds4: optional Metal output HC weights4 kernel unavailable\n");
        }

        MTLFunctionConstantValues *unary_sigmoid_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_sigmoid_op = 102;
        bool unary_cnt = false;
        [unary_sigmoid_constants setConstantValue:&unary_sigmoid_op type:MTLDataTypeShort atIndex:1200];
        [unary_sigmoid_constants setConstantValue:&unary_cnt        type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
                           constantValues:unary_sigmoid_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_sigmoid_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_sigmoid_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *unary_silu_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_silu_op = 106;
        [unary_silu_constants setConstantValue:&unary_silu_op type:MTLDataTypeShort atIndex:1200];
        [unary_silu_constants setConstantValue:&unary_cnt     type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
                           constantValues:unary_silu_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_silu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_silu_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *unary_softplus_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_softplus_op = 115;
        [unary_softplus_constants setConstantValue:&unary_softplus_op type:MTLDataTypeShort atIndex:1200];
        [unary_softplus_constants setConstantValue:&unary_cnt         type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
                           constantValues:unary_softplus_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_softplus_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_softplus_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *unary_sqrt_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_sqrt_op = 14;
        [unary_sqrt_constants setConstantValue:&unary_sqrt_op type:MTLDataTypeShort atIndex:1200];
        [unary_sqrt_constants setConstantValue:&unary_cnt     type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
                           constantValues:unary_sqrt_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_sqrt_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_sqrt_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *unary_clamp_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_clamp_op = 12;
        [unary_clamp_constants setConstantValue:&unary_clamp_op type:MTLDataTypeShort atIndex:1200];
        [unary_clamp_constants setConstantValue:&unary_cnt      type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32"
                           constantValues:unary_clamp_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_clamp_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_clamp_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *unary_scale_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_scale_op = 10;
        [unary_scale_constants setConstantValue:&unary_scale_op type:MTLDataTypeShort atIndex:1200];
        [unary_scale_constants setConstantValue:&unary_cnt      type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
                           constantValues:unary_scale_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_scale_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_scale_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        MTLFunctionConstantValues *unary_fill_constants = [[MTLFunctionConstantValues alloc] init];
        int16_t unary_fill_op = 11;
        [unary_fill_constants setConstantValue:&unary_fill_op type:MTLDataTypeShort atIndex:1200];
        [unary_fill_constants setConstantValue:&unary_cnt     type:MTLDataTypeBool  atIndex:1201];

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
                           constantValues:unary_fill_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_fill_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_fill_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        error = nil;
        fn = [library newFunctionWithName:@"kernel_unary_f16_f16"
                           constantValues:unary_fill_constants
                                    error:&error];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill function not found: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_unary_fill_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_unary_fill_f16_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        fn = [library newFunctionWithName:@"kernel_dsv4_hc_expand"];
        if (!fn) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand function not found\n");
            g_queue = nil;
            g_device = nil;
            return 0;
        }
        g_hc_expand_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
        if (!g_hc_expand_pipeline) {
            fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand pipeline failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_dsv4_indexer_score_one_direct_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct");
        g_dsv4_compressor_store_one_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one");
        g_dsv4_compressor_pack_ratio4_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4");
        g_dsv4_softmax_pool_ratio4_direct_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct");
        g_rms_norm_scale_pipeline =
            ds4_gpu_get_pipeline("kernel_rms_norm_scale_f32_4");
        g_dsv4_sort_i32_rows_asc_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_sort_i32_rows_asc");
        g_dsv4_indexed_attention_heads8_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8");
        g_dsv4_indexed_attention_heads8_rb16_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16");
        g_dsv4_softplus_sqrt_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4");
        g_dsv4_router_finalize_one_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one");
        g_dsv4_router_finalize_one_simd_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd");
        g_dsv4_router_finalize_weights_one_simd_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd");
        g_dsv4_router_weights_one_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one");
        g_glm_router_select_one_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_router_select_one");
        g_glm_kv_lora_rms_norm_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_kv_lora_rms_norm");
        g_glm_k_b_project_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_k_b_project_q8_0");
        g_glm_store_compact_kv_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_store_compact_kv");
        g_glm_qkv_norm_store_compact_kv_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_qkv_norm_store_compact_kv");
        g_glm_store_indexer_k_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_store_indexer_k");
        g_glm_build_kv_cache_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_build_kv_cache");
        g_glm_build_kv_cache_decode_group4_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_decode_group4");
        g_glm_build_kv_cache_flash_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_flash");
        g_glm_attention_full_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_full");
        g_glm_fill_selected_range_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_fill_selected_range");
        g_glm_fill_selected_range_batch_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_fill_selected_range_batch");
        g_glm_indexer_rope_tail_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_indexer_rope_tail_f32");
        g_glm_indexer_score_one_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_indexer_score_one");
        g_glm_indexer_score_one_direct_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct");
        g_glm_indexer_scores_batch_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch");
        g_glm_indexer_scores_tiled_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled");
        g_glm_indexer_scores_tiled_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled_f32");
        g_glm_qk_lowrank_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0");
        g_glm_qk_lowrank_glm52_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52");
        g_glm_qk_lowrank_glm52_sg_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52_sg");
        g_glm_qk_lowrank_batch_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch");
        g_glm_qk_lowrank_batch_glm52_t4_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_glm52_t4");
        g_glm_value_project_q8_0_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0");
        g_glm_value_project_q8_0_batch_heads_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads");
        g_glm_value_project_q8_0_batch_heads_mma_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads_mma");
        g_glm_attention_indexed_decode_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode");
        g_glm_attention_indexed_decode_split_group8_partial_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial");
        g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads");
        g_glm_attention_indexed_decode_split_group8_reduce_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce");
        g_glm_attention_indexed_decode_split_group8_reduce16_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce16");
        g_glm_attention_indexed_batch_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch");
        g_glm_attention_indexed_batch_group2_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group2");
        g_glm_attention_indexed_batch_q2_group4_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_q2_group4");
        g_glm_attention_indexed_batch_group8_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group8");
        g_glm_attention_indexed_batch_lora_group8_vec_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec");
        g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid");
        g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads");
        g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal");
        g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads");
        g_glm_q4_k_pair_swiglu_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu_f32");
        g_glm_q4_k_pair_swiglu2_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_f32");
        g_glm_q4_k_pair_swiglu4_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu4_f32");
        g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_f32");
        g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_row_f32");
        g_glm_q2_k_pair_swiglu_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q2_K_pair_swiglu_f32");
        g_glm_q2_k_addr_pair_swiglu2_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32");
        g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32_masked");
        g_glm_q4_k_addr_pair_swiglu_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32");
        g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32_masked");
        g_glm_q2_k_down_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q2_K_down_f32");
        g_glm_q4_k_down_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_down_simd_f32");
        g_glm_q2_k_addr_down_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_down_f32");
        g_glm_q4_k_addr_down_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_down_simd_f32");
        g_glm_q5_k_pair_swiglu_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_f32");
        g_glm_q5_k_pair_swiglu_mapped_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_f32");
        g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_row_f32");
        g_glm_q5_k_down_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q5_K_down_f32");
        g_glm_q6_k_down_f32_pipeline =
            ds4_gpu_get_pipeline("kernel_glm_q6_K_down_f32");
        g_dsv4_router_weights_batch_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch");
        g_dsv4_hc_expand4_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4");
        if (!g_dsv4_indexer_score_one_direct_pipeline ||
            !g_dsv4_compressor_store_one_pipeline ||
            !g_dsv4_sort_i32_rows_asc_pipeline ||
            !g_dsv4_indexed_attention_heads8_pipeline ||
            !g_dsv4_indexed_attention_heads8_rb16_pipeline ||
            !g_dsv4_softplus_sqrt_pipeline ||
            !g_dsv4_router_finalize_one_pipeline ||
            !g_dsv4_router_weights_one_pipeline ||
            !g_glm_router_select_one_pipeline ||
            !g_glm_kv_lora_rms_norm_pipeline ||
            !g_glm_k_b_project_pipeline ||
            !g_glm_store_compact_kv_pipeline ||
            !g_glm_qkv_norm_store_compact_kv_pipeline ||
            !g_glm_store_indexer_k_pipeline ||
            !g_glm_build_kv_cache_pipeline ||
            !g_glm_build_kv_cache_decode_group4_pipeline ||
            !g_glm_build_kv_cache_flash_pipeline ||
            !g_glm_attention_full_pipeline ||
            !g_glm_fill_selected_range_pipeline ||
            !g_glm_fill_selected_range_batch_pipeline ||
            !g_glm_indexer_rope_tail_pipeline ||
            !g_glm_indexer_score_one_pipeline ||
            !g_glm_indexer_score_one_direct_pipeline ||
            !g_glm_indexer_scores_batch_pipeline ||
            !g_glm_indexer_scores_tiled_pipeline ||
            !g_glm_indexer_scores_tiled_f32_pipeline ||
            !g_glm_qk_lowrank_pipeline ||
            !g_glm_qk_lowrank_glm52_pipeline ||
            !g_glm_qk_lowrank_glm52_sg_pipeline ||
            !g_glm_qk_lowrank_batch_pipeline ||
            !g_glm_qk_lowrank_batch_glm52_t4_pipeline ||
            !g_glm_value_project_q8_0_pipeline ||
            !g_glm_value_project_q8_0_batch_heads_pipeline ||
            !g_glm_value_project_q8_0_batch_heads_mma_pipeline ||
            !g_glm_attention_indexed_decode_pipeline ||
            !g_glm_attention_indexed_decode_split_group8_partial_pipeline ||
            !g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline ||
            !g_glm_attention_indexed_decode_split_group8_reduce_pipeline ||
            !g_glm_attention_indexed_decode_split_group8_reduce16_pipeline ||
            !g_glm_attention_indexed_batch_pipeline ||
            !g_glm_attention_indexed_batch_group2_pipeline ||
            !g_glm_attention_indexed_batch_q2_group4_pipeline ||
            !g_glm_attention_indexed_batch_group8_pipeline ||
            !g_glm_attention_indexed_batch_lora_group8_vec_pipeline ||
            !g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline ||
            !g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline ||
            !g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline ||
            !g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline ||
            !g_glm_q4_k_pair_swiglu_f32_pipeline ||
            !g_glm_q4_k_pair_swiglu2_f32_pipeline ||
            !g_glm_q4_k_pair_swiglu4_f32_pipeline ||
            !g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline ||
            !g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline ||
            !g_glm_q2_k_pair_swiglu_f32_pipeline ||
            !g_glm_q2_k_addr_pair_swiglu2_f32_pipeline ||
            !g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline ||
            !g_glm_q4_k_addr_pair_swiglu_f32_pipeline ||
            !g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline ||
            !g_glm_q2_k_down_f32_pipeline ||
            !g_glm_q4_k_down_f32_pipeline ||
            !g_glm_q2_k_addr_down_f32_pipeline ||
            !g_glm_q4_k_addr_down_f32_pipeline ||
            !g_glm_q5_k_pair_swiglu_f32_pipeline ||
            !g_glm_q5_k_pair_swiglu_mapped_f32_pipeline ||
            !g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline ||
            !g_glm_q5_k_down_f32_pipeline ||
            !g_glm_q6_k_down_f32_pipeline ||
            !g_dsv4_hc_expand4_pipeline) {
            g_queue = nil;
            g_device = nil;
            return 0;
        }

        g_initialized = 1;
    }

    return 1;
}

ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
    if (!g_initialized && !ds4_gpu_init()) return NULL;
    if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL;

    @autoreleasepool {
        DS4MetalTensor *tensor = [DS4MetalTensor new];
        tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes
                                              options:MTLResourceStorageModeShared];
        if (!tensor.buffer) {
            return NULL;
        }
        tensor.offset = 0;
        tensor.bytes = bytes;
        tensor.owner = 1;
        uint64_t live_snap = 0;
        uint64_t peak_snap = 0;
        pthread_mutex_lock(&g_tensor_mu);
        const int tracked = ds4_gpu_tensor_track_alloc_locked(
                (__bridge const void *)tensor,
                bytes,
                &live_snap,
                &peak_snap);
        pthread_mutex_unlock(&g_tensor_mu);
        if (!tracked) {
            fprintf(stderr, "ds4: failed to track Metal tensor allocation\n");
            tensor.buffer = nil;
            return NULL;
        }
        if (ds4_gpu_trace_allocs()) {
            fprintf(stderr,
                    "ds4: Metal tensor alloc %.3f MiB live %.3f MiB peak %.3f MiB\n",
                    (double)bytes / (1024.0 * 1024.0),
                    (double)live_snap / (1024.0 * 1024.0),
                    (double)peak_snap / (1024.0 * 1024.0));
        }
        return (__bridge_retained ds4_gpu_tensor *)tensor;
    }
}

ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) {
    return ds4_gpu_tensor_alloc(bytes);
}

int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) {
    (void)kv_cache_bytes;
    (void)context_bytes;
    return 0;
}

ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) {
    if (!base) return NULL;
    const DS4MetalTensor *base_obj = ds4_gpu_tensor_const_obj(base);
    if (offset > base_obj.bytes || bytes > base_obj.bytes - offset) return NULL;
    if (base_obj.offset > UINT64_MAX - offset) return NULL;
    const uint64_t absolute_offset = base_obj.offset + offset;
    if (absolute_offset > (uint64_t)NSUIntegerMax) return NULL;

    @autoreleasepool {
        DS4MetalTensor *view = [DS4MetalTensor new];
        view.buffer = base_obj.buffer;
        view.offset = absolute_offset;
        view.bytes = bytes;
        view.owner = 0;
        pthread_mutex_lock(&g_tensor_mu);
        const int tracked = ds4_gpu_tensor_track_view_locked((__bridge const void *)view);
        pthread_mutex_unlock(&g_tensor_mu);
        if (!tracked) {
            fprintf(stderr, "ds4: failed to track Metal tensor view\n");
            view.buffer = nil;
            return NULL;
        }
        return (__bridge_retained ds4_gpu_tensor *)view;
    }
}

void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) {
    if (!tensor) return;
    @autoreleasepool {
        uint8_t owner = 0;
        uint64_t bytes = 0;
        uint64_t live_snap = 0;
        uint64_t peak_snap = 0;
        if (!ds4_gpu_tensor_prepare_free(tensor,
                                         &owner,
                                         &bytes,
                                         &live_snap,
                                         &peak_snap)) {
            return;
        }
        DS4MetalTensor *obj = (__bridge_transfer DS4MetalTensor *)tensor;
        if (owner) {
            if (ds4_gpu_trace_allocs()) {
                fprintf(stderr,
                        "ds4: Metal tensor free %.3f MiB live %.3f MiB peak %.3f MiB\n",
                        (double)bytes / (1024.0 * 1024.0),
                        (double)live_snap / (1024.0 * 1024.0),
                        (double)peak_snap / (1024.0 * 1024.0));
            }
        }
        obj.buffer = nil;
        obj.offset = 0;
        obj.bytes = 0;
        obj.owner = 0;
    }
}

uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) {
    if (!tensor) return 0;
    const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
    return obj.bytes;
}

void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) {
    if (!tensor) return NULL;
    DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor);
    return (uint8_t *)[obj.buffer contents] + obj.offset;
}

int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) {
    if (!tensor || count > ds4_gpu_tensor_bytes(tensor) / sizeof(float)) return 0;
    float *p = ds4_gpu_tensor_contents(tensor);
    if (!p && count != 0) return 0;
    for (uint64_t i = 0; i < count; i++) p[i] = value;
    return 1;
}

int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) {
    if (!tensor || (!data && bytes != 0)) return 0;
    DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor);
    if (offset > obj.bytes || bytes > obj.bytes - offset) return 0;
    if (bytes != 0) {
        memcpy((uint8_t *)[obj.buffer contents] + obj.offset + offset, data, (size_t)bytes);
    }
    return 1;
}

int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) {
    if (!tensor || (!data && bytes != 0)) return 0;
    const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
    if (offset > obj.bytes || bytes > obj.bytes - offset) return 0;
    if (bytes != 0) {
        memcpy(data, (const uint8_t *)[obj.buffer contents] + obj.offset + offset, (size_t)bytes);
    }
    return 1;
}

int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset,
                          const ds4_gpu_tensor *src, uint64_t src_offset,
                          uint64_t bytes) {
    if (!dst || !src) return 0;
    if (!g_initialized && !ds4_gpu_init()) return 0;
    DS4MetalTensor *d = ds4_gpu_tensor_obj(dst);
    const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src);
    if (dst_offset > d.bytes || bytes > d.bytes - dst_offset) return 0;
    if (src_offset > s.bytes || bytes > s.bytes - src_offset) return 0;
    if (bytes == 0) return 1;
    if (!g_batch_cb) return 0;

    ds4_gpu_close_batch_encoder();
    g_batch_has_work = YES;
    id<MTLBlitCommandEncoder> blit = [g_batch_cb blitCommandEncoder];
    if (!blit) return 0;
    [blit copyFromBuffer:s.buffer
            sourceOffset:(NSUInteger)(s.offset + src_offset)
                toBuffer:d.buffer
       destinationOffset:(NSUInteger)(d.offset + dst_offset)
                    size:(NSUInteger)bytes];
    [blit endEncoding];
    return 1;
}

int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset,
                                   const ds4_gpu_tensor *src, uint64_t src_offset,
                                   uint64_t count) {
    if (!dst || !src) return 0;
    if (!g_initialized && !ds4_gpu_init()) return 0;
    DS4MetalTensor *d = ds4_gpu_tensor_obj(dst);
    const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src);
    if (count == 0) return 1;
    if (count > UINT64_MAX / sizeof(float) ||
        count > UINT64_MAX / sizeof(uint16_t)) {
        return 0;
    }
    const uint64_t src_bytes = count * sizeof(float);
    const uint64_t dst_bytes = count * sizeof(uint16_t);
    if (src_offset > s.bytes || src_bytes > s.bytes - src_offset ||
        dst_offset > d.bytes || dst_bytes > d.bytes - dst_offset) {
        return 0;
    }

    @autoreleasepool {
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        uint64_t done = 0;
        int ok = 1;
        while (done < count && ok) {
            uint64_t chunk64 = count - done;
            if (chunk64 > UINT32_MAX) chunk64 = UINT32_MAX;
            const uint32_t chunk = (uint32_t)chunk64;
            ok = ds4_gpu_encode_cpy_f32_f16_1d(
                    cb,
                    s.buffer,
                    (NSUInteger)(s.offset + src_offset + done * sizeof(float)),
                    d.buffer,
                    (NSUInteger)(d.offset + dst_offset + done * sizeof(uint16_t)),
                    chunk);
            done += chunk;
        }
        if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "tensor f32 to f16 copy");
        return ok;
    }
}

int ds4_gpu_pack_slot_rows_f32_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *slots,
        uint32_t                n_rows,
        uint32_t                width,
        uint32_t                n_slots,
        uint32_t                slot_cap) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 ||
        slot_cap == 0 || n_rows > slot_cap) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> slotsbuf = ds4_gpu_tensor_buffer(slots);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        uint64_t row_bytes = 0;
        uint64_t slot_plane_bytes = 0;
        uint64_t slots_bytes = 0;
        uint64_t out_rows = 0;
        uint64_t out_bytes = 0;
        if ((uint64_t)width > UINT64_MAX / sizeof(float)) return 0;
        row_bytes = (uint64_t)width * sizeof(float);
        if ((uint64_t)slot_cap > UINT64_MAX / row_bytes) return 0;
        slot_plane_bytes = (uint64_t)slot_cap * row_bytes;
        if ((uint64_t)n_slots > UINT64_MAX / slot_plane_bytes) return 0;
        slots_bytes = (uint64_t)n_slots * slot_plane_bytes;
        if ((uint64_t)n_rows > UINT64_MAX / n_slots) return 0;
        out_rows = (uint64_t)n_rows * n_slots;
        if (out_rows > UINT64_MAX / row_bytes) return 0;
        out_bytes = out_rows * row_bytes;
        if (!slotsbuf || !outbuf ||
            ds4_gpu_tensor_bytes(slots) < slots_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal slot-row pack received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;
        if (!ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb,
                                                        slotsbuf,
                                                        ds4_gpu_tensor_offset(slots),
                                                        outbuf,
                                                        ds4_gpu_tensor_offset(out),
                                                        width,
                                                        n_slots,
                                                        n_rows,
                                                        sizeof(float),
                                                        slot_plane_bytes,
                                                        row_bytes,
                                                        row_bytes,
                                                        (uint64_t)n_slots * row_bytes)) {
            return 0;
        }
        if (!ds4_gpu_finish_command_buffer(cb, owned, "slot-row pack")) return 0;
    }

    return 1;
}

int ds4_gpu_begin_commands(void) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (g_batch_cb) return 0;
    g_batch_cb = ds4_gpu_new_command_buffer();
    g_batch_has_work = NO;
    if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
    return g_batch_cb != nil;
}

int ds4_gpu_flush_encoder(void) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!g_batch_cb) return 0;
    ds4_gpu_close_batch_encoder();
    return 1;
}

int ds4_gpu_flush_commands(void) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!g_batch_cb) return 0;

    ds4_gpu_close_batch_encoder();
    id<MTLCommandBuffer> cb = g_batch_cb;
    g_batch_cb = nil;
    g_batch_has_work = NO;
    [cb commit];
    [g_pending_cbs addObject:cb];
    ds4_gpu_stream_expert_cache_note_batch_committed();

    g_batch_cb = ds4_gpu_new_command_buffer();
    g_batch_has_work = NO;
    if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
    if (!g_batch_cb) {
        (void)ds4_gpu_wait_pending_command_buffers("command batch");
        [g_transient_buffers removeAllObjects];
        return 0;
    }
    return 1;
}

int ds4_gpu_commands_active(void) {
    return g_batch_cb != nil;
}

static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label) {
    const char *what = label ? label : "streaming expert cache in-flight";
    if (g_batch_cb && ds4_gpu_flush_commands() == 0) return 0;
    if ([g_pending_cbs count] != 0 &&
        ds4_gpu_wait_pending_command_buffers(what) == 0) {
        return 0;
    }
    return 1;
}

int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) {
    if (!event_value) return 0;
    *event_value = 0;
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!g_batch_cb) return 0;

    if (@available(macOS 12.0, *)) {
        if (!g_selected_readback_event) {
            g_selected_readback_event = [g_device newSharedEvent];
            if (!g_selected_readback_event) {
                fprintf(stderr, "ds4: failed to create Metal shared event for selected-id overlap\n");
                return 0;
            }
        }

        ds4_gpu_close_batch_encoder();
        const uint64_t value = ++g_selected_readback_event_value;
        [g_batch_cb encodeSignalEvent:g_selected_readback_event value:value];
        g_batch_has_work = YES;
        *event_value = value;
        return 1;
    }

    fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n");
    return 0;
}

int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!g_batch_cb || event_value == 0) return 0;

    if (@available(macOS 12.0, *)) {
        if (!g_selected_readback_event) return 0;

        ds4_gpu_close_batch_encoder();
        id<MTLCommandBuffer> cb = g_batch_cb;
        g_batch_cb = nil;
        g_batch_has_work = NO;
        [cb commit];
        ds4_gpu_stream_expert_cache_note_batch_committed();

        const char *what = label ? label : "selected-id overlap";
        const BOOL signaled =
            [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000];
        [g_pending_cbs addObject:cb];
        if (!signaled) {
            fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what);
            (void)ds4_gpu_wait_pending_command_buffers(what);
            [g_transient_buffers removeAllObjects];
            return 0;
        }
        if (cb.status == MTLCommandBufferStatusError) {
            fprintf(stderr, "ds4: Metal %s failed: %s\n",
                    what,
                    [[cb.error localizedDescription] UTF8String]);
            (void)ds4_gpu_wait_pending_command_buffers(what);
            [g_transient_buffers removeAllObjects];
            return 0;
        }

        g_batch_cb = ds4_gpu_new_command_buffer();
        g_batch_has_work = NO;
        if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
        if (!g_batch_cb) {
            (void)ds4_gpu_wait_pending_command_buffers(what);
            [g_transient_buffers removeAllObjects];
            return 0;
        }
        return 1;
    }

    fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n");
    return 0;
}

/*
 * Tensor-parallel gates.
 *
 * A TP gate is a mid-command-stream rendezvous with the peer machine: the
 * kernels ahead of the gate leave a partial block output in a slab slot,
 * the GPU signals g_tp_gpu_event, and the pre-encoded combine kernel waits
 * on g_tp_cpu_event.  A dedicated service thread bridges the two: it spins
 * until the GPU reaches the gate, runs the transport exchange (RDMA WRITE
 * plus flag poll, or a TCP write/read pair — behind the callback), and
 * CPU-signals the release.  On exchange failure the release is signaled
 * anyway so the GPU never deadlocks; the failure latches in g_tp_failed
 * and the eval aborts at the next command-buffer boundary.
 *
 * Gate sequence values increase monotonically per encoded gate.  Both ranks
 * encode the identical graph, so the values agree by construction and slots
 * never need resetting between tokens.
 */
typedef struct {
    uint32_t layer;
    uint32_t gate;
    uint32_t rows;  /* 0 = row gate; >0 = verify-block batch gate */
    uint32_t event_arrival;
    uint64_t seq;
    /* Big batch gates (prefill): exchange big_bytes from big_out into
     * big_in directly (CPU-visible bounce buffers), bypassing the slab. */
    const void *big_out;
    void *big_in;
    uint64_t big_bytes;
} ds4_gpu_tp_request;

enum { DS4_GPU_TP_QUEUE = 1024 };

static id<MTLSharedEvent> g_tp_gpu_event;   /* GPU -> service thread */
static id<MTLSharedEvent> g_tp_cpu_event;   /* service thread -> GPU */
/* Batch (verify-block) gates run on their own sequence space and release
 * event: the row-gate seq feeds the RDMA pre-posted recv accounting, which
 * requires consecutive values, and a shared release event would make a
 * small batch value satisfy waits armed against the larger row seq. */
static id<MTLSharedEvent> g_tp_batch_gpu_event;
static id<MTLSharedEvent> g_tp_batch_cpu_event;
static uint64_t g_tp_batch_seq;
/* Batch flag values are tagged so a stale row-gate seq in the reused FFN
 * flag word can never satisfy a batch arrival spin (and vice versa). */
#define DS4_TP_BATCH_FLAG_TAG 0x80000000u
/* Expert-ownership split parameters for routed kernels. World 1 means TP is
 * not bound; world 2 assigns each rank one contiguous expert range. */
static int32_t g_tp_split_rank;
static int32_t g_tp_split_world = 1;
static int32_t g_tp_session_batch_mode;

static int ds4_gpu_tp_world_is_two(void) {
    return g_tp_split_world == 2;
}

/* Return the contiguous routed-expert range backed by this process. Rank 1
 * owns the high range and receives any odd-count remainder. */
static void ds4_gpu_tp_expert_range(uint32_t n_total_expert,
                                    uint32_t *first_expert,
                                    uint32_t *n_expert) {
    *first_expert = 0;
    *n_expert = n_total_expert;
    if (g_tp_split_world != 2) return;

    const uint32_t low_experts = n_total_expert / 2u;
    if (g_tp_split_rank == 1) {
        *first_expert = low_experts;
        *n_expert = n_total_expert - low_experts;
    } else {
        *n_expert = low_experts;
    }
}

/* Attention head split for GLM batch prefill: each rank computes a
 * contiguous half of the heads in the qk-low / attention-lora /
 * value-project batch kernels; the caller zeroes the unowned head range
 * of the heads buffer and combines the attn-output partials over the
 * TP big-gate exchange. */
static int32_t g_tp_attn_head_split;

void ds4_gpu_tp_set_attn_head_split(int enabled) {
    g_tp_attn_head_split = enabled ? 1 : 0;
}

static void ds4_gpu_tp_attn_head_range(uint32_t n_head,
                                       uint32_t group,
                                       uint32_t *head_base,
                                       uint32_t *head_count) {
    *head_base = 0;
    *head_count = n_head;
    if (!g_tp_attn_head_split || g_tp_split_world != 2) return;
    const uint32_t half = n_head / 2u;
    if (half == 0u || (half % group) != 0u || (n_head % 2u) != 0u) return;
    *head_count = half;
    *head_base = g_tp_split_rank == 1 ? half : 0u;
}
/* Flag gates (DS4_TP_FLAG_GATES): the GPU publishes gate arrival by storing
 * the sequence number into a slab word instead of signaling the shared
 * event; the service thread spin-reads it from shared memory, which wakes
 * hundreds of microseconds earlier than signaledValue polling.  The
 * CPU->GPU release direction stays on the shared event. */
static bool g_tp_flag_gates;
static id<MTLBuffer> g_tp_slab_buffer;
static NSUInteger g_tp_slab_buffer_off;
static volatile uint32_t *g_tp_gpu_flags;   /* CPU view of the flag words */
static uint64_t g_tp_gpu_flags_off;
static uint64_t g_tp_seq;
static ds4_gpu_tp_exchange_fn g_tp_exchange_fn;
static ds4_gpu_tp_batch_exchange_fn g_tp_batch_exchange_fn;
static ds4_gpu_tp_big_exchange_fn g_tp_big_exchange_fn;

void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) {
    g_tp_big_exchange_fn = fn;
}

static void *g_tp_exchange_ud;
static pthread_t g_tp_thread;
static int g_tp_thread_running;
static int g_tp_shutdown;
static int g_tp_failed_flag;
static pthread_mutex_t g_tp_mutex = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t g_tp_cond = PTHREAD_COND_INITIALIZER;
static ds4_gpu_tp_request g_tp_queue[DS4_GPU_TP_QUEUE];
static uint32_t g_tp_queue_head;
static uint32_t g_tp_queue_count;

static uint64_t g_tp_stat_gates;
static double g_tp_stat_gpu_wait_ms;
static double g_tp_stat_exchange_ms;

/* GPU keep-alive (see kernel_dsv4_tp_keepalive): its own queue and thread,
 * alive exactly as long as the TP gate machinery. */
static id<MTLCommandQueue> g_tp_keepalive_queue;
static id<MTLBuffer> g_tp_keepalive_buffer;
static pthread_t g_tp_keepalive_thread;
static int g_tp_keepalive_running;

/* Nonzero while a verify block runs: the GPU is genuinely busy
 * there, so the keep-alive is a pure parasite (~2.3ms per 5-row block
 * measured against the single-machine verify). */
static volatile int g_tp_keepalive_paused;

void ds4_gpu_tp_keepalive_pause(int paused) {
    g_tp_keepalive_paused = paused;
}

void ds4_gpu_tp_set_session_batch_mode(int enabled) {
    g_tp_session_batch_mode = enabled ? 1 : 0;
}

static uint32_t ds4_gpu_tp_keepalive_tgs_from_env(void) {
    uint32_t ka_tgs = 1;
    const char *tgs_env = getenv("DS4_TP_KEEPALIVE_TGS");
    if (tgs_env) {
        int v = atoi(tgs_env);
        if (v > 0 && v <= 2048) ka_tgs = (uint32_t)v;
    }
    return ka_tgs;
}

static void *ds4_gpu_tp_keepalive_thread(void *arg) {
    (void)arg;
    /* Swept on the M5 Max pair: too few iterations lets clocks sag.  Current
     * TP split-resident Flash runs show 1.2M is a small Q4/Q2 decode win over
     * 800k, while two threadgroups waste work. */
    uint32_t iters = 1200000;
    const char *env = getenv("DS4_TP_KEEPALIVE_ITERS");
    if (env) iters = (uint32_t)atoi(env);
    /* One ALU-only threadgroup keeps the GPU from power-gating but does
     * not push the frequency governor; solo-vs-engine kernel gaps
     * (~1.7x) suggest decode runs well below max clocks. More TGs raise
     * apparent utilization without eating memory bandwidth. */
    uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env();
    id<MTLComputePipelineState> pipeline =
        ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive");
    if (!pipeline) {
        fprintf(stderr, "ds4: TP keep-alive pipeline missing\n");
        return NULL;
    }
    while (!g_tp_shutdown) {
        if (g_tp_keepalive_paused) {
            usleep(200);
            continue;
        }
        @autoreleasepool {
            id<MTLCommandBuffer> cb = [g_tp_keepalive_queue commandBuffer];
            id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
            [enc setComputePipelineState:pipeline];
            [enc setBuffer:g_tp_keepalive_buffer offset:0 atIndex:0];
            [enc setBytes:&iters length:sizeof(iters) atIndex:1];
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)ka_tgs, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
            [enc endEncoding];
            [cb commit];
            [cb waitUntilCompleted];
        }
    }
    return NULL;
}

static void *ds4_gpu_tp_service_thread(void *arg) {
    (void)arg;
    const bool profile = getenv("DS4_TP_GATE_PROFILE") != NULL;
    while (1) {
        pthread_mutex_lock(&g_tp_mutex);
        while (g_tp_queue_count == 0 && !g_tp_shutdown)
            pthread_cond_wait(&g_tp_cond, &g_tp_mutex);
        if (g_tp_shutdown && g_tp_queue_count == 0) {
            pthread_mutex_unlock(&g_tp_mutex);
            break;
        }
        ds4_gpu_tp_request req = g_tp_queue[g_tp_queue_head];
        g_tp_queue_head = (g_tp_queue_head + 1) % DS4_GPU_TP_QUEUE;
        g_tp_queue_count--;
        pthread_mutex_unlock(&g_tp_mutex);

        /* Wait for the GPU to reach this gate.  Tight spin: gate arrival is
         * on the decode critical path and normally tens of microseconds
         * out; yielding here measurably delays the release wake-up. */
        const double t0 = profile ? ds4_gpu_now_ms() : 0.0;
        uint32_t spins = 0;
        if (req.big_bytes > 0) {
            /* Big gates always signal arrival through the batch shared
             * event (see ds4_gpu_tp_big_gate_kick): the event completion
             * semantics are what guarantee the bounce payload is visible
             * before the exchange reads it. */
            while (g_tp_batch_gpu_event.signaledValue < req.seq) {
                if (g_tp_shutdown) break;
                if (++spins > (1u << 16)) sched_yield();
            }
        } else if (!req.event_arrival) {
            const uint32_t slot = req.layer * 2u + req.gate;
            uint32_t want = (uint32_t)req.seq;
            if (req.rows > 0)
                want = DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq;
            while (__atomic_load_n(&g_tp_gpu_flags[slot], __ATOMIC_ACQUIRE) != want) {
                if (g_tp_shutdown) break;
                if (++spins > (1u << 20)) {
                    sched_yield();
                    spins = 0;
                }
            }
        } else if (req.rows > 0) {
            while (g_tp_batch_gpu_event.signaledValue < req.seq) {
                if (g_tp_shutdown) break;
                if (++spins > (1u << 16)) sched_yield();
            }
        } else {
            while (g_tp_gpu_event.signaledValue < req.seq) {
                if (g_tp_shutdown) break;
                if (++spins > (1u << 16)) sched_yield();
            }
        }
        const double t1 = profile ? ds4_gpu_now_ms() : 0.0;
        int ok = 0;
        if (!g_tp_shutdown && !g_tp_failed_flag) {
            if (req.big_bytes > 0) {
                if (g_tp_big_exchange_fn)
                    ok = g_tp_big_exchange_fn(g_tp_exchange_ud, req.layer,
                                              req.seq, req.big_out,
                                              req.big_in, req.big_bytes);
            } else if (req.rows > 0) {
                if (g_tp_batch_exchange_fn)
                    ok = g_tp_batch_exchange_fn(g_tp_exchange_ud, req.layer,
                                                req.rows, req.seq);
            } else if (g_tp_exchange_fn) {
                ok = g_tp_exchange_fn(g_tp_exchange_ud, req.layer, req.gate,
                                      req.seq);
            }
        }
        if (!ok && !g_tp_shutdown) {
            if (!g_tp_failed_flag)
                fprintf(stderr, "ds4: TP gate exchange failed (layer %u gate %u seq %llu)\n",
                        req.layer, req.gate, (unsigned long long)req.seq);
            g_tp_failed_flag = 1;
        }
        /* Release the GPU even on failure so end_commands can drain. */
        if (req.rows > 0) g_tp_batch_cpu_event.signaledValue = req.seq;
        else g_tp_cpu_event.signaledValue = req.seq;
        if (profile) {
            g_tp_stat_gpu_wait_ms += t1 - t0;
            g_tp_stat_exchange_ms += ds4_gpu_now_ms() - t1;
            if (++g_tp_stat_gates % 860 == 0) {
                fprintf(stderr,
                        "ds4: TP gates %llu: avg gpu-wait %.1f us, avg exchange %.1f us\n",
                        (unsigned long long)g_tp_stat_gates,
                        g_tp_stat_gpu_wait_ms / (double)g_tp_stat_gates * 1000.0,
                        g_tp_stat_exchange_ms / (double)g_tp_stat_gates * 1000.0);
            }
        }
    }
    return NULL;
}

int ds4_gpu_tp_init(uint32_t rank,
                    ds4_gpu_tensor *slab, uint64_t gpu_flags_off,
                    ds4_gpu_tp_exchange_fn fn, void *ud) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (g_tp_thread_running || rank > 1) return 0;
    g_tp_split_rank = (int32_t)rank;
    g_tp_split_world = 2;
    g_tp_slab_buffer = slab ? ds4_gpu_tensor_buffer(slab) : nil;
    g_tp_slab_buffer_off = slab ? ds4_gpu_tensor_offset(slab) : 0;
    g_tp_gpu_flags_off = gpu_flags_off;
    g_tp_gpu_flags = slab ?
        (volatile uint32_t *)((uint8_t *)ds4_gpu_tensor_contents(slab) + gpu_flags_off) : NULL;
    /* Flag arrival is the default: the slab-word publish detects in ~1-3us
     * where signaledValue polling costs 10-20, worth +2.3 t/s on the pair
     * (A/B 2026-07-06, byte-identical output).  DS4_TP_EVENT_GATES falls
     * back to the shared-event arrival path. */
    g_tp_flag_gates = g_tp_gpu_flags != NULL && getenv("DS4_TP_EVENT_GATES") == NULL;
    g_tp_gpu_event = [g_device newSharedEvent];
    g_tp_cpu_event = [g_device newSharedEvent];
    g_tp_batch_gpu_event = [g_device newSharedEvent];
    g_tp_batch_cpu_event = [g_device newSharedEvent];
    if (!g_tp_gpu_event || !g_tp_cpu_event ||
        !g_tp_batch_gpu_event || !g_tp_batch_cpu_event) {
        fprintf(stderr, "ds4: failed to create TP shared events\n");
        return 0;
    }
    g_tp_exchange_fn = fn;
    g_tp_exchange_ud = ud;
    g_tp_seq = 0;
    g_tp_batch_seq = 0;
    g_tp_shutdown = 0;
    g_tp_failed_flag = 0;
    g_tp_queue_head = 0;
    g_tp_queue_count = 0;
    pthread_attr_t attr;
    pthread_attr_init(&attr);
    pthread_attr_set_qos_class_np(&attr, QOS_CLASS_USER_INTERACTIVE, 0);
    if (pthread_create(&g_tp_thread, &attr, ds4_gpu_tp_service_thread, NULL) != 0) {
        pthread_attr_destroy(&attr);
        fprintf(stderr, "ds4: failed to start TP gate service thread\n");
        return 0;
    }
    pthread_attr_destroy(&attr);
    g_tp_thread_running = 1;
    if (getenv("DS4_TP_NO_KEEPALIVE") == NULL) {
        uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env();
        g_tp_keepalive_queue = [g_device newCommandQueue];
        g_tp_keepalive_buffer = [g_device newBufferWithLength:(NSUInteger)ka_tgs * 256u * sizeof(float)
                                                      options:MTLResourceStorageModeShared];
        if (g_tp_keepalive_queue && g_tp_keepalive_buffer &&
            pthread_create(&g_tp_keepalive_thread, NULL,
                           ds4_gpu_tp_keepalive_thread, NULL) == 0) {
            g_tp_keepalive_running = 1;
        } else {
            fprintf(stderr, "ds4: TP keep-alive setup failed (continuing without)\n");
        }
    }
    return 1;
}

void ds4_gpu_tp_shutdown(void) {
    if (!g_tp_thread_running) return;
    pthread_mutex_lock(&g_tp_mutex);
    g_tp_shutdown = 1;
    pthread_cond_broadcast(&g_tp_cond);
    pthread_mutex_unlock(&g_tp_mutex);
    pthread_join(g_tp_thread, NULL);
    g_tp_thread_running = 0;
    if (g_tp_keepalive_running) {
        pthread_join(g_tp_keepalive_thread, NULL);
        g_tp_keepalive_running = 0;
        g_tp_keepalive_queue = nil;
        g_tp_keepalive_buffer = nil;
    }
    g_tp_exchange_fn = NULL;
    g_tp_batch_exchange_fn = NULL;
    g_tp_exchange_ud = NULL;
    g_tp_split_rank = 0;
    g_tp_split_world = 1;
    g_tp_session_batch_mode = 0;
}

void ds4_gpu_tp_suspend_expert_sharding(int suspend) {
    if (!g_tp_thread_running) return;
    g_tp_split_world = suspend ? 1 : 2;
}

int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) {
    if (!g_batch_cb) {
        fprintf(stderr, "ds4: TP gate encode without an open command batch (layer %u gate %u)\n",
                layer, gate);
        return 0;
    }
    if (!g_tp_thread_running) {
        fprintf(stderr, "ds4: TP gate encode without the gate service (layer %u)\n", layer);
        return 0;
    }
    const uint64_t seq = ++g_tp_seq;
    const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates;
    if (!event_arrival) {
        /* Publish arrival through the slab word; the buffer hazard against
         * the partial-output kernels orders the store after the payload. */
        const uint32_t slot = layer * 2u + gate;
        const uint32_t value = (uint32_t)seq;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) return 0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set");
        if (!pipeline) return 0;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBuffer:g_tp_slab_buffer
                offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u)
               atIndex:0];
        [enc setBytes:&value length:sizeof(value) atIndex:1];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        ds4_gpu_close_batch_encoder();
    } else {
        ds4_gpu_close_batch_encoder();
        [g_batch_cb encodeSignalEvent:g_tp_gpu_event value:seq];
    }
    [g_batch_cb encodeWaitForEvent:g_tp_cpu_event value:seq];
    pthread_mutex_lock(&g_tp_mutex);
    if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) {
        pthread_mutex_unlock(&g_tp_mutex);
        fprintf(stderr, "ds4: TP gate queue overflow\n");
        return 0;
    }
    uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE;
    g_tp_queue[tail].layer = layer;
    g_tp_queue[tail].gate = gate;
    g_tp_queue[tail].rows = 0;
    g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u;
    g_tp_queue[tail].seq = seq;
    g_tp_queue[tail].big_out = NULL;
    g_tp_queue[tail].big_in = NULL;
    g_tp_queue[tail].big_bytes = 0;
    g_tp_queue_count++;
    pthread_cond_signal(&g_tp_cond);
    pthread_mutex_unlock(&g_tp_mutex);
    return 1;
}

void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) {
    g_tp_batch_exchange_fn = fn;
}

/* Verify-block batch gate: same arrival/release machinery as the row gate
 * (the FFN flag word and event pair are reused — a decode gate and a batch
 * gate are never in flight together, and seq values stay globally unique),
 * but the service thread runs the multi-row exchange callback. */
int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) {
    if (!g_batch_cb) return 0;
    if (!g_tp_thread_running || rows == 0) return 0;
    const uint64_t seq = ++g_tp_batch_seq;
    const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates;
    if (!event_arrival) {
        const uint32_t slot = layer * 2u + 1u; /* FFN gate slot */
        const uint32_t value = DS4_TP_BATCH_FLAG_TAG | (uint32_t)seq;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) return 0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set");
        if (!pipeline) return 0;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBuffer:g_tp_slab_buffer
                offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u)
               atIndex:0];
        [enc setBytes:&value length:sizeof(value) atIndex:1];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        ds4_gpu_close_batch_encoder();
    } else {
        ds4_gpu_close_batch_encoder();
        [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq];
    }
    [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq];
    pthread_mutex_lock(&g_tp_mutex);
    if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) {
        pthread_mutex_unlock(&g_tp_mutex);
        fprintf(stderr, "ds4: TP gate queue overflow\n");
        return 0;
    }
    uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE;
    g_tp_queue[tail].layer = layer;
    g_tp_queue[tail].gate = 1u; /* FFN */
    g_tp_queue[tail].rows = rows;
    g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u;
    g_tp_queue[tail].seq = seq;
    g_tp_queue[tail].big_out = NULL;
    g_tp_queue[tail].big_in = NULL;
    g_tp_queue[tail].big_bytes = 0;
    g_tp_queue_count++;
    pthread_cond_signal(&g_tp_cond);
    pthread_mutex_unlock(&g_tp_mutex);
    return 1;
}

/* Prefill batch gate kick: same seq space and release event as the verify
 * batch gate, but the service thread exchanges big_bytes directly between
 * the two shared bounce buffers instead of slab slots.  The kick only
 * publishes the GPU arrival marker and queues the exchange; the caller
 * encodes the release wait later through ds4_gpu_tp_big_gate_wait, which
 * lets it interleave more GPU work with the wire exchange.  Arrival always
 * uses the batch shared event, NOT the flag word: a flag write carries no
 * memory-visibility guarantee for the payload buffer, and once the GPU
 * keeps running past the kick (no event wait right behind it) the service
 * thread can observe the flag before the producing kernels' stores reach
 * CPU-visible memory (measured: stale rows in the first sub-kick).  The
 * shared-event signal only fires after every preceding command completes,
 * which is exactly the payload ordering the exchange needs; the ~10 us
 * slower arrival detection is noise against a multi-ms exchange. */
uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows,
                                  const ds4_gpu_tensor *out_t,
                                  ds4_gpu_tensor *in_t,
                                  uint64_t bytes) {
    if (!g_batch_cb) return 0;
    if (!g_tp_thread_running || rows == 0 || bytes == 0) return 0;
    const void *out_ptr = ds4_gpu_tensor_contents((ds4_gpu_tensor *)out_t);
    void *in_ptr = ds4_gpu_tensor_contents(in_t);
    if (!out_ptr || !in_ptr) {
        fprintf(stderr, "ds4: TP big gate needs CPU-visible bounce buffers\n");
        return 0;
    }
    const uint64_t seq = ++g_tp_batch_seq;
    ds4_gpu_close_batch_encoder();
    [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq];
    pthread_mutex_lock(&g_tp_mutex);
    if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) {
        pthread_mutex_unlock(&g_tp_mutex);
        fprintf(stderr, "ds4: TP gate queue overflow\n");
        return 0;
    }
    uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE;
    g_tp_queue[tail].layer = layer;
    g_tp_queue[tail].gate = 1u;
    g_tp_queue[tail].rows = rows;
    g_tp_queue[tail].event_arrival = 1u;
    g_tp_queue[tail].seq = seq;
    g_tp_queue[tail].big_out = out_ptr;
    g_tp_queue[tail].big_in = in_ptr;
    g_tp_queue[tail].big_bytes = bytes;
    g_tp_queue_count++;
    pthread_cond_signal(&g_tp_cond);
    pthread_mutex_unlock(&g_tp_mutex);
    return seq;
}

/* Encode the GPU-side release wait for a previously kicked big gate.  The
 * batch release event is monotonic and the service thread completes queued
 * exchanges in kick order, so waiting on the LAST kicked seq of a stage
 * also covers every earlier kick. */
int ds4_gpu_tp_big_gate_wait(uint64_t seq) {
    if (!g_batch_cb || seq == 0) return 0;
    ds4_gpu_close_batch_encoder();
    [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq];
    return 1;
}

int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows,
                               const ds4_gpu_tensor *out_t,
                               ds4_gpu_tensor *in_t,
                               uint64_t bytes) {
    const uint64_t seq = ds4_gpu_tp_big_gate_kick(layer, rows, out_t, in_t, bytes);
    if (seq == 0) return 0;
    return ds4_gpu_tp_big_gate_wait(seq);
}

int ds4_gpu_tp_failed(void) {
    return g_tp_failed_flag;
}

int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (event_value == 0) return 0;

    if (@available(macOS 12.0, *)) {
        if (!g_selected_readback_event) return 0;

        const char *what = label ? label : "selected-id readback";
        const BOOL signaled =
            [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000];
        if (!signaled) {
            fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what);
            return 0;
        }
        return 1;
    }

    fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n");
    return 0;
}

static int ds4_gpu_signal_batch_and_wait_event(const char *label) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!g_batch_cb) return 0;

    if (@available(macOS 12.0, *)) {
        if (!g_selected_readback_event) {
            g_selected_readback_event = [g_device newSharedEvent];
            if (!g_selected_readback_event) {
                fprintf(stderr, "ds4: failed to create Metal shared event for %s\n",
                        label ? label : "selected-id readback");
                return 0;
            }
        }

        ds4_gpu_close_batch_encoder();
        id<MTLCommandBuffer> cb = g_batch_cb;
        g_batch_cb = nil;
        const uint64_t value = ++g_selected_readback_event_value;
        [cb encodeSignalEvent:g_selected_readback_event value:value];
        g_batch_has_work = YES;
        [cb commit];
        ds4_gpu_stream_expert_cache_note_batch_committed();

        const BOOL signaled = [g_selected_readback_event waitUntilSignaledValue:value timeoutMS:60000];
        [g_pending_cbs addObject:cb];
        if (!signaled) {
            fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n",
                    label ? label : "selected-id readback");
            (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback");
            [g_transient_buffers removeAllObjects];
            return 0;
        }
        if (cb.status == MTLCommandBufferStatusError) {
            fprintf(stderr, "ds4: Metal %s failed: %s\n",
                    label ? label : "selected-id readback",
                    [[cb.error localizedDescription] UTF8String]);
            (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback");
            [g_transient_buffers removeAllObjects];
            return 0;
        }

        g_batch_cb = ds4_gpu_new_command_buffer();
        g_batch_has_work = NO;
        if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
        if (!g_batch_cb) {
            (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback");
            [g_transient_buffers removeAllObjects];
            return 0;
        }
        return 1;
    } else {
        ds4_gpu_close_batch_encoder();
        id<MTLCommandBuffer> cb = g_batch_cb;
        g_batch_cb = nil;
        g_batch_has_work = NO;
        if (ds4_gpu_finish_command_buffer(cb, 1, label ? label : "selected-id readback") == 0) {
            return 0;
        }
        return ds4_gpu_begin_commands();
    }
}

int ds4_gpu_end_commands(void) {
    if (!g_batch_cb) return 0;
    ds4_gpu_close_batch_encoder();
    id<MTLCommandBuffer> cb = g_batch_cb;
    g_batch_cb = nil;
    g_batch_has_work = NO;
    g_stream_expert_cache_owned_seq = g_stream_expert_cache_batch_seq;
    g_stream_expert_cache_batch_seq = 0;
    return ds4_gpu_finish_command_buffer(cb, 1, "command batch");
}

static int ds4_gpu_flash_attn_stage_profile_boundary(
        id<MTLCommandBuffer> __strong *cbp,
        const char           *mode,
        const char           *stage,
        uint32_t              n_tokens,
        uint32_t              n_comp,
        uint32_t              n_keys,
        uint32_t              n_head,
        uint32_t              head_dim,
        uint32_t              window,
        uint32_t              ratio,
        double               *stage_t0) {
    if (!cbp || !*cbp || !stage_t0 || !stage) return 0;
    if (ds4_gpu_end_commands() == 0) return 0;

    const double now_ms = ds4_gpu_now_ms();
    const char *filter = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE_FILTER");
    const int print_stage =
        !filter || !filter[0] ||
        strstr(stage, filter) != NULL ||
        (mode && strstr(mode, filter) != NULL);
    if (print_stage) {
        fprintf(stderr,
                "ds4: Metal FlashAttention prefill stage mode=%s tokens=%u comp=%u "
                "keys=%u heads=%u dim=%u window=%u ratio=%u %s=%.3f ms\n",
                mode ? mode : "unknown",
                n_tokens,
                n_comp,
                n_keys,
                n_head,
                head_dim,
                window,
                ratio,
                stage,
                now_ms - *stage_t0);
    }
    *stage_t0 = now_ms;

    if (ds4_gpu_begin_commands() == 0) return 0;
    int owned = 0;
    *cbp = ds4_gpu_command_buffer(&owned);
    return *cbp != nil && owned == 0;
}

int ds4_gpu_synchronize(void) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (g_batch_cb) return ds4_gpu_end_commands();
    if ([g_pending_cbs count] != 0) {
        int ok = ds4_gpu_wait_pending_command_buffers("synchronize");
        [g_transient_buffers removeAllObjects];
        ds4_gpu_model_buffer_cache_maybe_evict("synchronize");
        return ok;
    }

    id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
    if (!cb) return 0;
    return ds4_gpu_finish_command_buffer(cb, 1, "synchronize");
}

void ds4_gpu_cleanup(void) {
    if (!g_initialized) return;

    @autoreleasepool {
        if (g_batch_cb) {
            ds4_gpu_close_batch_encoder();
            [g_batch_cb commit];
            [g_batch_cb waitUntilCompleted];
            g_batch_cb = nil;
            if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_done_seq) {
                g_stream_expert_cache_done_seq = g_stream_expert_cache_batch_seq;
            }
            g_stream_expert_cache_batch_seq = 0;
        }
        (void)ds4_gpu_wait_pending_command_buffers("cleanup");
        if (ds4_gpu_stream_expert_timing_summary_enabled() &&
            getenv("DS4_METAL_MEMORY_REPORT") == NULL) {
            ds4_gpu_print_memory_report("at cleanup");
        }
        g_selected_readback_event = nil;
        g_selected_readback_event_value = 0;
        [g_transient_buffers removeAllObjects];
        ds4_gpu_stream_expert_pread_pool_shutdown();
        ds4_gpu_stream_expert_cache_clear_all(1);
        for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) {
            g_stream_expert_cache_gate_addr_buffers[layer] = nil;
            g_stream_expert_cache_up_addr_buffers[layer] = nil;
            g_stream_expert_cache_down_addr_buffers[layer] = nil;
            g_stream_compact_gate_addr_buffers[layer] = nil;
            g_stream_compact_up_addr_buffers[layer] = nil;
            g_stream_compact_down_addr_buffers[layer] = nil;
            g_stream_compact_selected_buffers[layer] = nil;
            g_stream_selected_id_buffers[layer] = nil;
        }
        g_set_rows_f32_i32_pipeline = nil;
        g_get_rows_f32_pipeline = nil;
        g_get_rows_f16_pipeline = nil;
        g_get_rows_i32_pipeline = nil;
        g_get_rows_q8_0_pipeline = nil;
        g_get_rows_q4_0_pipeline = nil;
        g_get_rows_q4_K_pipeline = nil;
        g_repeat_f32_pipeline = nil;
        g_concat_pipeline = nil;
        g_cpy_f32_f32_pipeline = nil;
        g_cpy_f32_f16_pipeline = nil;
        g_cpy_contig_f32_f16_pipeline = nil;
        g_cpy_f16_f32_pipeline = nil;
        g_cpy_f16_f16_pipeline = nil;
        g_cpy_contig_f16_f32_pipeline = nil;
        g_cpy_contig_f16_f16_pipeline = nil;
        g_flash_kv_stage_f16_pipeline = nil;
        g_swiglu_pipeline = nil;
        g_swiglu_flat_pipeline = nil;
        g_add_pipeline = nil;
        g_add2_pipeline = nil;
        g_add3_pipeline = nil;
        g_moe_sum6_pipeline = nil;
        g_moe_sum8_pipeline = nil;
        g_mul_pipeline = nil;
        g_bin_mul_scalar_pipeline = nil;
        g_bin_div_row_pipeline = nil;
        g_unary_sigmoid_pipeline = nil;
        g_unary_silu_pipeline = nil;
        g_unary_softplus_pipeline = nil;
        g_unary_sqrt_pipeline = nil;
        g_unary_clamp_pipeline = nil;
        g_unary_scale_pipeline = nil;
        g_unary_fill_pipeline = nil;
        g_unary_fill_f16_pipeline = nil;
        g_rms_norm_pipeline = nil;
        g_rms_norm_plain_pipeline = nil;
        g_add_rms_norm_pipeline = nil;
        g_rms_norm_scale_pipeline = nil;
        g_dsv4_qkv_rms_norm_pipeline = nil;
        g_hc_split_sinkhorn_pipeline = nil;
        g_hc_split_weighted_sum_pipeline = nil;
        g_hc_split_weighted_sum_norm_pipeline = nil;
        g_hc_weighted_sum_pipeline = nil;
        g_hc_weighted_sum_norm_pipeline = nil;
        g_output_hc_weights4_pipeline = nil;
        g_hc_expand_pipeline = nil;
        g_moe_mul_mv_id_iq2_xxs_pipeline = nil;
        g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil;
        g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_id_q2_k_pipeline = nil;
        g_moe_mul_mv_id_q2_k_sum6_pipeline = nil;
        g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = nil;
        g_moe_mul_mv_id_q4_k_pipeline = nil;
        g_moe_mul_mv_id_q4_k_pair_pipeline = nil;
        g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_id_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_group_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_group6_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_group8_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_group24_q4_k_id_pipeline = nil;
        g_moe_mul_mv_group24_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_slots6_q2_k_sum6_pipeline = nil;
        g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_slots6_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_addr_iq2_xxs_pipeline = nil;
        g_moe_mul_mv_addr_q2_k_sum6_pipeline = nil;
        g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = nil;
        g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = nil;
        g_moe_stream_expert_cache_validate_pipeline = nil;
        g_moe_q4_gather_slots6_pipeline = nil;
        g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_table_q4_k_sum6_pipeline = nil;
        g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = nil;
        g_moe_mul_mv_addr_q4_k_sum6_pipeline = nil;
        g_moe_table_q4_pair_gate_encoder = nil;
        g_moe_table_q4_pair_up_encoder = nil;
        g_moe_table_q4_sum_down_encoder = nil;
        g_rope_tail_batch_pipeline = nil;
        g_rope_tail_inplace_pair_pipeline = nil;
        g_rope_tail_inplace_pair_shared4_pipeline = nil;
        g_rope_tail_inplace_pair_affine_pipeline = nil;
        g_dsv4_fp8_kv_quantize_pipeline = nil;
        g_dsv4_indexer_qat_pipeline = nil;
        g_dsv4_kv_fp8_store_pipeline = nil;
        g_dsv4_ratio4_shift_pipeline = nil;
        g_dsv4_compressor_pack_ratio4_pipeline = nil;
        g_dsv4_softmax_pool_ratio4_direct_pipeline = nil;
        g_dsv4_softmax_pool_pipeline = nil;
        g_soft_max_f32_pipeline = nil;
        g_soft_max_f32_4_pipeline = nil;
        g_argsort_f32_i32_desc_pipeline = nil;
        g_argsort_merge_f32_i32_desc_pipeline = nil;
        g_sum_rows_f32_f32_pipeline = nil;
        g_dsv4_topk_mask_pipeline = nil;
        g_dsv4_topk_mask_scatter_pipeline = nil;
        g_dsv4_indexer_weighted_sum_pipeline = nil;
        g_dsv4_indexer_score_one_direct_pipeline = nil;
        g_dsv4_compressor_store_one_pipeline = nil;
        g_dsv4_sort_i32_rows_asc_pipeline = nil;
        g_dsv4_indexed_attention_heads8_pipeline = nil;
        g_dsv4_indexed_attention_heads8_rb16_pipeline = nil;
        g_dsv4_softplus_sqrt_pipeline = nil;
        g_dsv4_router_finalize_one_pipeline = nil;
        g_dsv4_router_finalize_one_simd_pipeline = nil;
        g_dsv4_router_finalize_weights_one_simd_pipeline = nil;
        g_dsv4_router_weights_one_pipeline = nil;
        g_glm_router_select_one_pipeline = nil;
        g_glm_kv_lora_rms_norm_pipeline = nil;
        g_glm_k_b_project_pipeline = nil;
        g_glm_store_compact_kv_pipeline = nil;
        g_glm_qkv_norm_store_compact_kv_pipeline = nil;
        g_glm_store_indexer_k_pipeline = nil;
        g_glm_build_kv_cache_pipeline = nil;
        g_glm_build_kv_cache_decode_group4_pipeline = nil;
        g_glm_build_kv_cache_flash_pipeline = nil;
        g_glm_attention_full_pipeline = nil;
        g_glm_fill_selected_range_pipeline = nil;
        g_glm_fill_selected_range_batch_pipeline = nil;
        g_glm_indexer_rope_tail_pipeline = nil;
        g_glm_indexer_score_one_pipeline = nil;
        g_glm_indexer_score_one_direct_pipeline = nil;
        g_glm_indexer_scores_batch_pipeline = nil;
        g_glm_indexer_scores_tiled_pipeline = nil;
        g_glm_indexer_scores_tiled_f32_pipeline = nil;
        g_glm_qk_lowrank_pipeline = nil;
        g_glm_qk_lowrank_glm52_pipeline = nil;
        g_glm_qk_lowrank_glm52_sg_pipeline = nil;
        g_glm_qk_lowrank_batch_pipeline = nil;
        g_glm_qk_lowrank_batch_glm52_t4_pipeline = nil;
        g_glm_value_project_q8_0_pipeline = nil;
        g_glm_value_project_q8_0_batch_heads_pipeline = nil;
        g_glm_value_project_q8_0_batch_heads_mma_pipeline = nil;
        g_glm_attention_indexed_decode_pipeline = nil;
        g_glm_attention_indexed_decode_split_group8_partial_pipeline = nil;
        g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = nil;
        g_glm_attention_indexed_decode_split_group8_reduce_pipeline = nil;
        g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = nil;
        g_glm_attention_indexed_batch_pipeline = nil;
        g_glm_attention_indexed_batch_group2_pipeline = nil;
        g_glm_attention_indexed_batch_q2_group4_pipeline = nil;
        g_glm_attention_indexed_batch_group8_pipeline = nil;
        g_glm_attention_indexed_batch_lora_group8_vec_pipeline = nil;
        g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = nil;
        g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = nil;
        g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = nil;
        g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = nil;
        g_glm_q4_k_pair_swiglu_f32_pipeline = nil;
        g_glm_q4_k_pair_swiglu2_f32_pipeline = nil;
        g_glm_q4_k_pair_swiglu4_f32_pipeline = nil;
        g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = nil;
        g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = nil;
        g_glm_q2_k_pair_swiglu_f32_pipeline = nil;
        g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = nil;
        g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = nil;
        g_glm_q4_k_addr_pair_swiglu_f32_pipeline = nil;
        g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = nil;
        g_glm_q2_k_down_f32_pipeline = nil;
        g_glm_q4_k_down_f32_pipeline = nil;
        g_glm_q2_k_addr_down_f32_pipeline = nil;
        g_glm_q4_k_addr_down_f32_pipeline = nil;
        g_glm_q5_k_pair_swiglu_f32_pipeline = nil;
        g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = nil;
        g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = nil;
        g_glm_q5_k_down_f32_pipeline = nil;
        g_glm_q6_k_down_f32_pipeline = nil;
        g_dsv4_router_weights_batch_pipeline = nil;
        g_dsv4_hc_expand4_pipeline = nil;
        g_flash_attn_mask_buffer = nil;
        g_flash_attn_zero_mask_buffer = nil;
        g_flash_attn_pad_buffer = nil;
        g_flash_attn_tmp_buffer = nil;
        g_flash_attn_blk_buffer = nil;
        ds4_gpu_clear_zero_prefix_prefill_mask_cache();
        g_flash_attn_ring_buffer = nil;
        g_flash_attn_kv_buffer = nil;
        g_glm_flash_attn_mask_buffer = nil;
        g_compressor_pool_kv_buffer = nil;
        g_compressor_pool_score_buffer = nil;
        g_compressor_pool_score_cont_buffer = nil;
        g_compressor_pool_softmax_buffer = nil;
        g_compressor_pool_product_buffer = nil;
        g_compressor_store_ape_buffer = nil;
        g_compressor_store_score_buffer = nil;
        g_embed_rows_buffer = nil;
        g_router_selection_buffer = nil;
        g_router_weight_sum_buffer = nil;
        g_indexer_head_scores_buffer = nil;
        g_indexer_topk_buffer = nil;
        g_indexed_topk_buffer = nil;
        g_stream_expert_validate_status_buffer = nil;
        g_f16_round_scratch_buffer = nil;
        g_raw_store_round_buffer = nil;
        g_moe_gate_scratch_buffer = nil;
        g_moe_down_scratch_buffer = nil;
        g_moe_id_map_buffer = nil;
        g_moe_q4_gate_slots_buffer = nil;
        g_moe_q4_up_slots_buffer = nil;
        g_moe_q4_down_slots_buffer = nil;
        g_attn_out_group_ids_buffer = nil;
        g_model_fd = -1;
        g_model_map_ptr = NULL;
        g_model_map_size = 0;
        g_model_mapped_offset = 0;
        g_model_mapped_size = 0;
        g_model_mapped_max_tensor_bytes = 0;
        ds4_gpu_tensor_tracking_reset();
        g_flash_attn_mask_bytes = 0;
        g_flash_attn_zero_mask_bytes = 0;
        g_flash_attn_pad_bytes = 0;
        g_flash_attn_tmp_bytes = 0;
        g_flash_attn_blk_bytes = 0;
        g_flash_attn_ring_bytes = 0;
        g_flash_attn_kv_bytes = 0;
        g_glm_flash_attn_mask_bytes = 0;
        g_glm_flash_attn_mask_valid = 0;
        g_glm_flash_attn_mask_pos0 = 0;
        g_glm_flash_attn_mask_tokens = 0;
        g_glm_flash_attn_mask_cache_len = 0;
        g_compressor_pool_kv_bytes = 0;
        g_compressor_pool_score_bytes = 0;
        g_compressor_pool_score_cont_bytes = 0;
        g_compressor_pool_softmax_bytes = 0;
        g_compressor_pool_product_bytes = 0;
        g_compressor_store_ape_bytes = 0;
        g_compressor_store_score_bytes = 0;
        g_embed_rows_bytes = 0;
        g_router_selection_bytes = 0;
        g_router_weight_sum_bytes = 0;
        g_indexer_head_scores_bytes = 0;
        g_indexer_topk_bytes = 0;
        g_indexed_topk_bytes = 0;
        g_f16_round_scratch_bytes = 0;
        g_raw_store_round_bytes = 0;
        g_moe_gate_scratch_bytes = 0;
        g_moe_down_scratch_bytes = 0;
        g_moe_id_map_bytes = 0;
        g_moe_q4_gate_slots_bytes = 0;
        g_moe_q4_up_slots_bytes = 0;
        g_moe_q4_down_slots_bytes = 0;
        g_attn_out_group_ids_bytes = 0;
        g_model_wrap_count = 0;
        g_model_wrap_bytes = 0;
        g_model_wrap_max_bytes = 0;
        g_model_buffer_cache_bytes = 0;
        g_model_buffer_cache_evictions = 0;
        g_model_buffer_cache_over_limit = 0;
        ds4_gpu_model_residency_clear();
        ds4_gpu_model_views_clear();
        [g_pipeline_cache removeAllObjects];
        g_pipeline_cache = nil;
        [g_q4_expert_layer_residency_cache removeAllObjects];
        g_q4_expert_layer_residency_cache = nil;
        [g_q4_expert_table_cache removeAllObjects];
        g_q4_expert_table_cache = nil;
        [g_model_buffer_cache removeAllObjects];
        g_model_buffer_cache = nil;
        g_transient_buffers = nil;
        g_pending_cbs = nil;
        g_library = nil;
        g_queue = nil;
        g_device = nil;
        g_initialized = 0;
    }
}

static int ds4_gpu_encode_get_rows_f16(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        weight,
        NSUInteger           weight_offset,
        id<MTLBuffer>        tokens,
        NSUInteger           tokens_offset,
        id<MTLBuffer>        out,
        NSUInteger           out_offset,
        uint32_t             n_vocab,
        uint32_t             n_tokens,
        uint32_t             n_embd) {
    if (!cb || !weight || !tokens || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
        return 0;
    }

    const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t);
    const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float);
    const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
    ds4_gpu_get_rows_args args = {
        .ne00t = (int32_t)n_embd,
        .ne00 = (int32_t)n_embd,
        .nb01 = src_row_bytes,
        .nb02 = (uint64_t)n_vocab * src_row_bytes,
        .nb03 = (uint64_t)n_vocab * src_row_bytes,
        .ne10 = (int32_t)n_tokens,
        .nb10 = sizeof(int32_t),
        .nb11 = token_bytes,
        .nb12 = token_bytes,
        .nb1 = dst_row_bytes,
        .nb2 = (uint64_t)n_tokens * dst_row_bytes,
        .nb3 = (uint64_t)n_tokens * dst_row_bytes,
    };

    NSUInteger nth = (NSUInteger)n_embd;
    const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1;
    const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_get_rows_f16_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:weight offset:weight_offset atIndex:1];
    [enc setBuffer:tokens offset:tokens_offset atIndex:2];
    [enc setBuffer:out offset:out_offset atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static uint64_t ds4_gpu_q8_0_row_bytes(uint32_t n_embd) {
    return (((uint64_t)n_embd + 31u) / 32u) * 34u;
}

static int ds4_gpu_quant_row_bytes(
        uint32_t  type,
        uint32_t  n_embd,
        uint64_t *row_bytes_out) {
    if (!row_bytes_out || n_embd == 0) return 0;
    switch (type) {
    case DS4_METAL_TENSOR_Q8_0:
        *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 34u;
        return 1;
    case DS4_METAL_TENSOR_Q4_0:
        *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 18u;
        return 1;
    case DS4_METAL_TENSOR_Q4_K:
        if ((n_embd % 256u) != 0) return 0;
        *row_bytes_out = ((uint64_t)n_embd / 256u) * 144u;
        return 1;
    default:
        return 0;
    }
}

static int ds4_gpu_q8_0_table_bytes(
        uint32_t  n_vocab,
        uint32_t  n_embd,
        uint64_t *bytes_out) {
    if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0;
    const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd);
    if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0;
    *bytes_out = (uint64_t)n_vocab * row_bytes;
    return 1;
}

static int ds4_gpu_quant_table_bytes(
        uint32_t  type,
        uint32_t  n_vocab,
        uint32_t  n_embd,
        uint64_t *bytes_out) {
    if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0;
    uint64_t row_bytes = 0;
    if (!ds4_gpu_quant_row_bytes(type, n_embd, &row_bytes)) return 0;
    if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0;
    *bytes_out = (uint64_t)n_vocab * row_bytes;
    return 1;
}

static int ds4_gpu_encode_get_rows_q8_0(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        weight,
        NSUInteger           weight_offset,
        id<MTLBuffer>        tokens,
        NSUInteger           tokens_offset,
        const int32_t       *single_token,
        id<MTLBuffer>        out,
        NSUInteger           out_offset,
        uint32_t             n_vocab,
        uint32_t             n_tokens,
        uint32_t             n_embd) {
    if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
        return 0;
    }
    if (!tokens && (!single_token || n_tokens != 1)) {
        return 0;
    }

    ds4_gpu_get_rows_q8_0_args args = {
        .n_embd = (int32_t)n_embd,
        .n_vocab = (int32_t)n_vocab,
        .n_tokens = (int32_t)n_tokens,
        .src_row_bytes = ds4_gpu_q8_0_row_bytes(n_embd),
        .dst_row_bytes = (uint64_t)n_embd * sizeof(float),
        .token_stride = sizeof(int32_t),
    };

    NSUInteger nth = 32u;
    const NSUInteger max_threads = g_get_rows_q8_0_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1;
    const NSUInteger nblocks = ((NSUInteger)n_embd + 31u) / 32u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_get_rows_q8_0_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:weight offset:weight_offset atIndex:1];
    if (tokens) {
        [enc setBuffer:tokens offset:tokens_offset atIndex:2];
    } else {
        [enc setBytes:single_token length:sizeof(*single_token) atIndex:2];
    }
    [enc setBuffer:out offset:out_offset atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_get_rows_quant(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        weight,
        NSUInteger           weight_offset,
        uint32_t             weight_type,
        id<MTLBuffer>        tokens,
        NSUInteger           tokens_offset,
        const int32_t       *single_token,
        id<MTLBuffer>        out,
        NSUInteger           out_offset,
        uint32_t             n_vocab,
        uint32_t             n_tokens,
        uint32_t             n_embd) {
    if (weight_type == DS4_METAL_TENSOR_Q8_0) {
        return ds4_gpu_encode_get_rows_q8_0(cb,
                                            weight,
                                            weight_offset,
                                            tokens,
                                            tokens_offset,
                                            single_token,
                                            out,
                                            out_offset,
                                            n_vocab,
                                            n_tokens,
                                            n_embd);
    }
    if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
        return 0;
    }
    if (!tokens && (!single_token || n_tokens != 1)) {
        return 0;
    }

    uint64_t src_row_bytes = 0;
    if (!ds4_gpu_quant_row_bytes(weight_type, n_embd, &src_row_bytes)) return 0;
    ds4_gpu_get_rows_q8_0_args args = {
        .n_embd = (int32_t)n_embd,
        .n_vocab = (int32_t)n_vocab,
        .n_tokens = (int32_t)n_tokens,
        .src_row_bytes = src_row_bytes,
        .dst_row_bytes = (uint64_t)n_embd * sizeof(float),
        .token_stride = sizeof(int32_t),
    };

    id<MTLComputePipelineState> pipeline = nil;
    NSUInteger block_width = 0;
    if (weight_type == DS4_METAL_TENSOR_Q4_0) {
        pipeline = g_get_rows_q4_0_pipeline;
        block_width = 32u;
    } else if (weight_type == DS4_METAL_TENSOR_Q4_K) {
        pipeline = g_get_rows_q4_K_pipeline;
        block_width = 256u;
    }
    if (!pipeline || block_width == 0) return 0;

    NSUInteger nth = block_width;
    const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1;
    const NSUInteger nblocks = ((NSUInteger)n_embd + block_width - 1u) / block_width;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:weight offset:weight_offset atIndex:1];
    if (tokens) {
        [enc setBuffer:tokens offset:tokens_offset atIndex:2];
    } else {
        [enc setBytes:single_token length:sizeof(*single_token) atIndex:2];
    }
    [enc setBuffer:out offset:out_offset atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_repeat_hc_embedding(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        rows,
        NSUInteger           rows_offset,
        id<MTLBuffer>        out,
        NSUInteger           out_offset,
        uint32_t             n_tokens,
        uint32_t             n_embd,
        uint32_t             n_hc) {
    if (!cb || !rows || !out || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0;

    const uint64_t embd_bytes = (uint64_t)n_embd * sizeof(float);
    ds4_gpu_repeat_args args = {
        .ne00 = (int32_t)n_embd,
        .ne01 = 1,
        .ne02 = (int32_t)n_tokens,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = embd_bytes,
        .nb02 = embd_bytes,
        .nb03 = (uint64_t)n_tokens * embd_bytes,
        .ne0 = (int32_t)n_embd,
        .ne1 = (int32_t)n_hc,
        .ne2 = (int32_t)n_tokens,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = embd_bytes,
        .nb2 = (uint64_t)n_hc * embd_bytes,
        .nb3 = (uint64_t)n_tokens * n_hc * embd_bytes,
    };

    NSUInteger nth = (NSUInteger)n_embd;
    const NSUInteger max_threads = g_repeat_f32_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_repeat_f32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:rows offset:rows_offset atIndex:1];
    [enc setBuffer:out offset:out_offset atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(n_hc, n_tokens, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

int ds4_gpu_embed_token_q8_0_tensor(
        ds4_gpu_tensor *out,
        const void       *model_map,
        uint64_t          model_size,
        uint64_t          weight_offset,
        uint32_t          n_vocab,
        uint32_t          token,
        uint32_t          n_embd) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float);
        if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal Q8_0 embedding received undersized output buffer\n");
            return 0;
        }

        uint64_t weight_bytes = 0;
        if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) ||
            weight_offset > model_size ||
            weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal Q8_0 embedding range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        uint32_t token_for_kernel = token;
        id<MTLBuffer> wbuf = nil;
        const bool exact_token_row =
            getenv("DS4_METAL_DISABLE_TOKEN_EMBED_EXACT_VIEW") == NULL;
        if (exact_token_row) {
            const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd);
            const uint64_t token_rel = (uint64_t)token * row_bytes;
            if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) {
                fprintf(stderr, "ds4: Metal Q8_0 embedding token row is outside the mapped table\n");
                return 0;
            }
            wbuf = ds4_gpu_wrap_model_exact_range(model_map,
                                                  model_size,
                                                  weight_offset + token_rel,
                                                  row_bytes,
                                                  &inner_offset);
            token_for_kernel = 0;
        } else {
            wbuf = ds4_gpu_wrap_model_range(model_map,
                                            model_size,
                                            weight_offset,
                                            weight_bytes,
                                            &inner_offset);
        }
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const int32_t token_i32 = (int32_t)token_for_kernel;
        if (!ds4_gpu_encode_get_rows_q8_0(cb,
                                           wbuf,
                                           (NSUInteger)inner_offset,
                                           nil,
                                           0,
                                           &token_i32,
                                           outbuf,
                                           ds4_gpu_tensor_offset(out),
                                           n_vocab,
                                           1,
                                           n_embd)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed token")) return 0;
    }

    return 1;
}

int ds4_gpu_embed_tokens_q8_0_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *tokens,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                n_vocab,
        uint32_t                n_tokens,
        uint32_t                n_embd) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> tokbuf = ds4_gpu_tensor_buffer(tokens);
        const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float);
        const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
        if (!outbuf || !tokbuf ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(tokens) < token_bytes) {
            fprintf(stderr, "ds4: Metal Q8_0 batched embedding received undersized buffers\n");
            return 0;
        }

        uint64_t weight_bytes = 0;
        if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) ||
            weight_offset > model_size ||
            weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal Q8_0 batched embedding range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_get_rows_q8_0(cb,
                                           wbuf,
                                           (NSUInteger)inner_offset,
                                           tokbuf,
                                           ds4_gpu_tensor_offset(tokens),
                                           NULL,
                                           outbuf,
                                           ds4_gpu_tensor_offset(out),
                                           n_vocab,
                                           n_tokens,
                                           n_embd)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed tokens")) return 0;
    }

    return 1;
}

int ds4_gpu_embed_token_quant_tensor(
        ds4_gpu_tensor *out,
        const void       *model_map,
        uint64_t          model_size,
        uint64_t          weight_offset,
        uint32_t          weight_type,
        uint32_t          n_vocab,
        uint32_t          token,
        uint32_t          n_embd) {
    if (weight_type == DS4_METAL_TENSOR_Q8_0) {
        return ds4_gpu_embed_token_q8_0_tensor(out,
                                               model_map,
                                               model_size,
                                               weight_offset,
                                               n_vocab,
                                               token,
                                               n_embd);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float);
        if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal quant embedding received undersized output buffer\n");
            return 0;
        }

        uint64_t weight_bytes = 0;
        uint64_t row_bytes = 0;
        if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) ||
            !ds4_gpu_quant_row_bytes(weight_type, n_embd, &row_bytes) ||
            weight_offset > model_size ||
            weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal quant embedding range is outside the mapped model\n");
            return 0;
        }

        const uint64_t token_rel = (uint64_t)token * row_bytes;
        if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) {
            fprintf(stderr, "ds4: Metal quant embedding token row is outside the mapped table\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           weight_offset + token_rel,
                                           row_bytes,
                                           &inner_offset);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const int32_t token_i32 = 0;
        if (!ds4_gpu_encode_get_rows_quant(cb,
                                           wbuf,
                                           (NSUInteger)inner_offset,
                                           weight_type,
                                           nil,
                                           0,
                                           &token_i32,
                                           outbuf,
                                           ds4_gpu_tensor_offset(out),
                                           n_vocab,
                                           1,
                                           n_embd)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed token")) return 0;
    }

    return 1;
}

int ds4_gpu_embed_tokens_quant_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *tokens,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                weight_type,
        uint32_t                n_vocab,
        uint32_t                n_tokens,
        uint32_t                n_embd) {
    if (weight_type == DS4_METAL_TENSOR_Q8_0) {
        return ds4_gpu_embed_tokens_q8_0_tensor(out,
                                                tokens,
                                                model_map,
                                                model_size,
                                                weight_offset,
                                                n_vocab,
                                                n_tokens,
                                                n_embd);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> tokbuf = ds4_gpu_tensor_buffer(tokens);
        const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float);
        const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
        if (!outbuf || !tokbuf ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(tokens) < token_bytes) {
            fprintf(stderr, "ds4: Metal quant batched embedding received undersized buffers\n");
            return 0;
        }

        uint64_t weight_bytes = 0;
        if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) ||
            weight_offset > model_size ||
            weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal quant batched embedding range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_get_rows_quant(cb,
                                           wbuf,
                                           (NSUInteger)inner_offset,
                                           weight_type,
                                           tokbuf,
                                           ds4_gpu_tensor_offset(tokens),
                                           NULL,
                                           outbuf,
                                           ds4_gpu_tensor_offset(out),
                                           n_vocab,
                                           n_tokens,
                                           n_embd)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed tokens")) return 0;
    }

    return 1;
}

int ds4_gpu_embed_token_hc_tensor(
        ds4_gpu_tensor *out_hc,
        const void       *model_map,
        uint64_t          model_size,
        uint64_t          weight_offset,
        uint32_t          n_vocab,
        uint32_t          token,
        uint32_t          n_embd,
        uint32_t          n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out_hc || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0 || n_hc == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
        const uint64_t out_bytes = (uint64_t)n_embd * n_hc * sizeof(float);
        if (!outbuf || ds4_gpu_tensor_bytes(out_hc) < out_bytes) {
            fprintf(stderr, "ds4: Metal graph embedding received undersized HC output buffer\n");
            return 0;
        }

        const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t);
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal graph embedding range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        const NSUInteger row_bytes = (NSUInteger)n_embd * sizeof(float);
        if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer,
                                             &g_embed_rows_bytes,
                                             row_bytes,
                                             "ds4_embed_rows")) {
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const int32_t token_i32 = (int32_t)token;
        const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t);
        const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float);
        ds4_gpu_get_rows_args args = {
            .ne00t = (int32_t)n_embd,
            .ne00 = (int32_t)n_embd,
            .nb01 = src_row_bytes,
            .nb02 = (uint64_t)n_vocab * src_row_bytes,
            .nb03 = (uint64_t)n_vocab * src_row_bytes,
            .ne10 = 1,
            .nb10 = sizeof(int32_t),
            .nb11 = sizeof(int32_t),
            .nb12 = sizeof(int32_t),
            .nb1 = dst_row_bytes,
            .nb2 = dst_row_bytes,
            .nb3 = dst_row_bytes,
        };
        NSUInteger nth = (NSUInteger)n_embd;
        const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > max_threads) nth = max_threads;
        if (nth == 0) nth = 1;
        const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_get_rows_f16_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
        [enc setBytes:&token_i32 length:sizeof(token_i32) atIndex:2];
        [enc setBuffer:g_embed_rows_buffer offset:0 atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake(nw0, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_encode_repeat_hc_embedding(cb,
                                                  g_embed_rows_buffer,
                                                  0,
                                                  outbuf,
                                                  ds4_gpu_tensor_offset(out_hc),
                                                  1,
                                                  n_embd,
                                                  n_hc)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed token")) return 0;
    }

    return 1;
}

int ds4_gpu_embed_tokens_hc_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *tokens,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                n_vocab,
        uint32_t                n_tokens,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out_hc || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0 || n_hc == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
        id<MTLBuffer> tokbuf = ds4_gpu_tensor_buffer(tokens);
        const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * n_hc * sizeof(float);
        const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
        if (!outbuf || !tokbuf ||
            ds4_gpu_tensor_bytes(out_hc) < out_bytes ||
            ds4_gpu_tensor_bytes(tokens) < token_bytes) {
            fprintf(stderr, "ds4: Metal graph batched embedding received undersized buffers\n");
            return 0;
        }

        const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t);
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal graph batched embedding range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        const NSUInteger rows_bytes = (NSUInteger)n_tokens * n_embd * sizeof(float);
        if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer,
                                             &g_embed_rows_bytes,
                                             rows_bytes,
                                             "ds4_embed_rows")) {
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_get_rows_f16(cb,
                                           wbuf,
                                           (NSUInteger)inner_offset,
                                           tokbuf,
                                           ds4_gpu_tensor_offset(tokens),
                                           g_embed_rows_buffer,
                                           0,
                                           n_vocab,
                                           n_tokens,
                                           n_embd) ||
            !ds4_gpu_encode_repeat_hc_embedding(cb,
                                                  g_embed_rows_buffer,
                                                  0,
                                                  outbuf,
                                                  ds4_gpu_tensor_offset(out_hc),
                                                  n_tokens,
                                                  n_embd,
                                                  n_hc)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed tokens")) return 0;
    }

    return 1;
}

int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!model_map || model_size == 0) return 0;
    if (map_offset > model_size || map_size == 0 || map_size > model_size - map_offset) return 0;
    max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size, max_tensor_bytes);

    @autoreleasepool {
        if (g_model_map_ptr == model_map &&
            g_model_map_size == model_size &&
            g_model_mapped_offset == map_offset &&
            g_model_mapped_size == map_size &&
            g_model_mapped_max_tensor_bytes == max_tensor_bytes) {
            return 1;
        }

        for (uint32_t i = 0; i < g_model_view_count; i++) {
            if (g_model_views[i].model_map == model_map &&
                g_model_views[i].model_size == model_size &&
                map_offset >= g_model_views[i].model_offset &&
                map_offset + map_size <= g_model_views[i].model_offset + g_model_views[i].bytes) {
                return 1;
            }
        }

        ds4_gpu_model_residency_clear();
        if (!ds4_gpu_map_model_views(model_map, model_size, map_offset, map_size, max_tensor_bytes)) {
            ds4_gpu_model_residency_clear();
            return 0;
        }
        g_model_map_ptr = model_map;
        g_model_map_size = model_size;
        g_model_mapped_offset = map_offset;
        g_model_mapped_size = map_size;
        g_model_mapped_max_tensor_bytes = max_tensor_bytes;
        if (ds4_gpu_model_map_log_enabled()) {
            fprintf(stderr,
                    "ds4: Metal mapped mmaped model as %u overlapping shared buffers\n",
                    g_model_view_count);
        }
        return 1;
    }
}

static int ds4_gpu_model_views_cover_spans(
        const void     *model_map,
        uint64_t        model_size,
        const uint64_t *offsets,
        const uint64_t *sizes,
        uint32_t        count);

int ds4_gpu_set_model_map_spans(
        const void *model_map,
        uint64_t model_size,
        const uint64_t *offsets,
        const uint64_t *sizes,
        uint32_t count,
        uint64_t max_tensor_bytes) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) return 0;
    if (count == 1) {
        return ds4_gpu_set_model_map_range(model_map,
                                           model_size,
                                           offsets[0],
                                           sizes[0],
                                           max_tensor_bytes);
    }
    if (ds4_gpu_model_views_cover_spans(model_map, model_size, offsets, sizes, count)) {
        return 1;
    }

    @autoreleasepool {
        const double t0 = ds4_gpu_now_ms();
        max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(model_size, max_tensor_bytes);

        ds4_gpu_model_residency_clear();
        ds4_gpu_model_views_clear();

        uint64_t mapped_total = 0;
        uint64_t first_offset = UINT64_MAX;
        for (uint32_t i = 0; i < count; i++) {
            if (offsets[i] > model_size || sizes[i] == 0 || sizes[i] > model_size - offsets[i]) {
                fprintf(stderr, "ds4: Metal model span %u is outside the GGUF mapping\n", i);
                ds4_gpu_model_residency_clear();
                ds4_gpu_model_views_clear();
                return 0;
            }
            if (offsets[i] < first_offset) first_offset = offsets[i];
            uint64_t effective_max = max_tensor_bytes;
            if (effective_max > sizes[i]) effective_max = sizes[i];
            if (!ds4_gpu_add_model_view_range(model_map,
                                              model_size,
                                              offsets[i],
                                              sizes[i],
                                              effective_max,
                                              true,
                                              &mapped_total)) {
                ds4_gpu_model_residency_clear();
                ds4_gpu_model_views_clear();
                return 0;
            }
        }
        if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) {
            ds4_gpu_model_residency_clear();
            ds4_gpu_model_views_clear();
            return 0;
        }
        g_model_map_ptr = model_map;
        g_model_map_size = model_size;
        g_model_mapped_offset = first_offset == UINT64_MAX ? 0 : first_offset;
        g_model_mapped_size = mapped_total;
        g_model_mapped_max_tensor_bytes = max_tensor_bytes;
        if (ds4_gpu_model_map_log_enabled()) {
            fprintf(stderr,
                    "ds4: Metal mapped mmaped model as %u disjoint shared buffers across %u tensor spans\n",
                    g_model_view_count,
                    count);
        }
        return 1;
    }
}

int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) {
    return ds4_gpu_set_model_map_range(model_map, model_size, 0, model_size, 0);
}

int ds4_gpu_set_model_fd(int fd) {
    g_model_fd = fd;
    return 1;
}

static int ds4_gpu_model_views_cover_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    size) {
    if (!model_map || model_size == 0 || size == 0 ||
        offset > model_size || size > model_size - offset) {
        return 0;
    }
    const uint64_t end = offset + size;
    for (uint32_t i = 0; i < g_model_view_count; i++) {
        if (g_model_views[i].model_map != model_map ||
            g_model_views[i].model_size != model_size) {
            continue;
        }
        const uint64_t view_start = g_model_views[i].model_offset;
        const uint64_t view_end = view_start + g_model_views[i].bytes;
        if (offset >= view_start && end <= view_end) return 1;
    }
    return 0;
}

static int ds4_gpu_model_views_cover_spans(
        const void     *model_map,
        uint64_t        model_size,
        const uint64_t *offsets,
        const uint64_t *sizes,
        uint32_t        count) {
    if (!offsets || !sizes || count == 0) return 0;
    for (uint32_t i = 0; i < count; i++) {
        if (!ds4_gpu_model_views_cover_range(model_map,
                                             model_size,
                                             offsets[i],
                                             sizes[i])) {
            return 0;
        }
    }
    return 1;
}

int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) {
    (void)fd;
    (void)model_map;
    return 1;
}

static id<MTLBuffer> ds4_gpu_wrap_model_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset) {
    (void)model_map;
    if (model_size == 0 || offset > model_size || len > model_size - offset) {
        fprintf(stderr, "ds4: Metal model range is outside the mapped model\n");
        return nil;
    }

    const uint64_t end = offset + len;
    for (uint32_t i = 0; i < g_model_view_count; i++) {
        if (g_model_views[i].model_map != model_map ||
            g_model_views[i].model_size != model_size) {
            continue;
        }
        const uint64_t view_start = g_model_views[i].model_offset;
        const uint64_t view_end = view_start + g_model_views[i].bytes;
        if (offset >= view_start && end <= view_end) {
            *inner_offset = offset - view_start;
            return g_model_views[i].buffer;
        }
    }

    fprintf(stderr,
            "ds4: Metal model range %.2f..%.2f GiB is not covered by mapped model views\n",
            ds4_gpu_gib(offset),
            ds4_gpu_gib(end));
    return nil;
}

typedef enum {
    DS4_GPU_EXACT_VIEW_CACHED,
    DS4_GPU_EXACT_VIEW_TRANSIENT,
    DS4_GPU_EXACT_VIEW_OWNED,
} ds4_gpu_exact_view_lifetime;

static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_impl(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset,
        ds4_gpu_exact_view_lifetime lifetime) {
    const bool cache_view = lifetime == DS4_GPU_EXACT_VIEW_CACHED;
    const bool transient_view = lifetime == DS4_GPU_EXACT_VIEW_TRANSIENT;
    if (!model_map || !g_device ||
        (cache_view && !g_model_buffer_cache) ||
        (transient_view && !g_transient_buffers) ||
        model_size == 0 || offset > model_size || len > model_size - offset) {
        fprintf(stderr, "ds4: Metal exact model range is outside the mapped model\n");
        return nil;
    }

    const uint64_t page = (uint64_t)getpagesize();
    const uint64_t page_offset = offset & ~(page - 1);
    const uint64_t leading = offset - page_offset;
    if (len > UINT64_MAX - leading ||
        leading + len > UINT64_MAX - (page - 1)) {
        fprintf(stderr, "ds4: Metal exact model range overflows page alignment\n");
        return nil;
    }
    uint64_t view_bytes = round_up_u64(leading + len, page);
    if (view_bytes > model_size - page_offset) view_bytes = model_size - page_offset;
    if (leading + len > view_bytes) {
        fprintf(stderr, "ds4: Metal exact model range alignment exceeds mapped model\n");
        return nil;
    }
    if (view_bytes > (uint64_t)[g_device maxBufferLength]) {
        fprintf(stderr,
                "ds4: Metal exact model range %.2f GiB exceeds maxBufferLength %.2f GiB\n",
                ds4_gpu_gib(view_bytes),
                ds4_gpu_gib((uint64_t)[g_device maxBufferLength]));
        return nil;
    }

    NSString *key = nil;
    id<MTLBuffer> buffer = nil;
    if (cache_view) {
        key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu",
               model_map,
               (unsigned long long)model_size,
               (unsigned long long)page_offset,
               (unsigned long long)view_bytes];
        buffer = [g_model_buffer_cache objectForKey:key];
    }
    if (!buffer) {
        const uintptr_t base = (uintptr_t)model_map;
        buffer = [g_device newBufferWithBytesNoCopy:(void *)(base + page_offset)
                                             length:(NSUInteger)view_bytes
                                            options:ds4_gpu_model_resource_options()
                                        deallocator:nil];
        if (!buffer) {
            fprintf(stderr,
                    "ds4: Metal could not wrap exact mmaped model range at %.2f GiB, size %.2f MiB\n",
                    ds4_gpu_gib(page_offset),
                    ds4_gpu_mib(view_bytes));
            return nil;
        }
        if (cache_view) {
            buffer.label = @"ds4_model_exact_view";
        } else if (transient_view) {
            buffer.label = @"ds4_model_exact_transient_view";
        } else {
            buffer.label = @"ds4_model_exact_owned_view";
        }
        if (cache_view) {
            [g_model_buffer_cache setObject:buffer forKey:key];
            ds4_gpu_model_buffer_cache_note_insert(view_bytes);
        } else if (transient_view) {
            [g_transient_buffers addObject:buffer];
        }
    }

    if (inner_offset) *inner_offset = leading;
    return buffer;
}

static id<MTLBuffer> ds4_gpu_wrap_model_exact_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset) {
    return ds4_gpu_wrap_model_exact_range_impl(model_map,
                                               model_size,
                                               offset,
                                               len,
                                               inner_offset,
                                               DS4_GPU_EXACT_VIEW_CACHED);
}

static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_transient(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset) {
    return ds4_gpu_wrap_model_exact_range_impl(model_map,
                                               model_size,
                                               offset,
                                               len,
                                               inner_offset,
                                               DS4_GPU_EXACT_VIEW_TRANSIENT);
}

static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_owned(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t   *inner_offset) {
    return ds4_gpu_wrap_model_exact_range_impl(model_map,
                                               model_size,
                                               offset,
                                               len,
                                               inner_offset,
                                               DS4_GPU_EXACT_VIEW_OWNED);
}

static id<MTLBuffer> ds4_gpu_wrap_q8_decode_model_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t    n_tokens,
        uint64_t   *inner_offset) {
    const uint64_t exact_decode_max_mib =
        ds4_gpu_env_u64("DS4_METAL_Q8_DECODE_EXACT_VIEW_MAX_MIB",
                        1024u,
                        1u,
                        4096u);
    const uint64_t exact_decode_max_bytes =
        exact_decode_max_mib * 1024ull * 1024ull;
    const bool exact_decode_weight_view =
        n_tokens == 1u &&
        len <= exact_decode_max_bytes &&
        getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL &&
        getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL;
    return exact_decode_weight_view ?
        ds4_gpu_wrap_model_exact_range(model_map,
                                       model_size,
                                       offset,
                                       len,
                                       inner_offset) :
        ds4_gpu_wrap_model_range(model_map,
                                 model_size,
                                 offset,
                                 len,
                                 inner_offset);
}

static id<MTLBuffer> ds4_gpu_wrap_f32_decode_model_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len,
        uint64_t    n_tokens,
        uint64_t   *inner_offset) {
    const uint64_t exact_decode_max_mib =
        ds4_gpu_env_u64("DS4_METAL_F32_DECODE_EXACT_VIEW_MAX_MIB",
                        64u,
                        1u,
                        4096u);
    const uint64_t exact_decode_max_bytes =
        exact_decode_max_mib * 1024ull * 1024ull;
    const bool exact_decode_weight_view =
        n_tokens == 1u &&
        len <= exact_decode_max_bytes &&
        getenv("DS4_METAL_ENABLE_F32_DECODE_EXACT_VIEWS") != NULL &&
        getenv("DS4_METAL_DISABLE_F32_DECODE_EXACT_VIEWS") == NULL;
    return exact_decode_weight_view ?
        ds4_gpu_wrap_model_exact_range(model_map,
                                       model_size,
                                       offset,
                                       len,
                                       inner_offset) :
        ds4_gpu_wrap_model_range(model_map,
                                 model_size,
                                 offset,
                                 len,
                                 inner_offset);
}

uint32_t ds4_gpu_stream_expert_cache_configured_count(void) {
    uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
    if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
        budget = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES;
    }
    return budget;
}

uint32_t ds4_gpu_stream_expert_cache_current_count(void) {
    return g_stream_expert_cache_entry_count;
}

uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes) {
    if (!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes)) {
        return 0;
    }
    return ds4_gpu_stream_expert_cache_configured_budget();
}

static int ds4_gpu_stream_expert_cache_note_expert_size(
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes) {
    if (gate_expert_bytes == 0 || down_expert_bytes == 0) return 0;
    if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
        fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n");
        return 0;
    }
    /*
     * The cache is a single-size-class slab allocator: the expert byte size is
     * frozen on first sight (or pre-seeded at startup from the model's slab
     * class) and off-size layers are REJECTED rather than adopted. A rejected
     * layer (mixed-precision boost: Q4_K experts among IQ2 layers) falls back
     * to the mapped-model per-expert path; last-writer-wins here would instead
     * poison the slab size class and deadlock slab reuse.
     */
    const uint64_t bytes = gate_expert_bytes * 2ull + down_expert_bytes;
    if (g_stream_expert_cache_expert_bytes == 0) {
        g_stream_expert_cache_expert_bytes = bytes;
        return 1;
    }
    return bytes == g_stream_expert_cache_expert_bytes;
}

static uint32_t ds4_gpu_stream_expert_cache_requested_budget(void) {
    if (!g_ssd_streaming_mode) return 0;
    if (g_stream_expert_cache_budget_override != 0) {
        return g_stream_expert_cache_budget_override;
    }
    return 0;
}

static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void) {
    return ds4_gpu_stream_expert_cache_requested_budget();
}

static uint32_t ds4_gpu_stream_expert_cache_effective_cap(
        uint32_t layer,
        uint32_t n_total_expert,
        uint32_t n_selected) {
    (void)layer;
    if (ds4_gpu_stream_expert_cache_configured_budget() == 0) return 0;

    /*
     * The residency policy is global: every layer can use any expert slot it
     * routes to, and global pruning decides which existing entry is least
     * valuable.  A per-layer cap made cache size depend on model depth rather
     * than the actual byte budget.
     */
    uint32_t cap = n_total_expert;
    if (cap < n_selected) cap = n_selected;
    if (cap > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
        cap = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
    }
    return cap;
}

static int ds4_gpu_stream_expert_timing_summary_enabled(void) {
    static int checked = 0;
    static int enabled = 0;
    if (!checked) {
        enabled =
            (getenv("DS4_METAL_STREAMING_EXPERT_TIMING_SUMMARY") != NULL ||
             getenv("DS4_METAL_STREAMING_EXPERT_PROFILE_SUMMARY") != NULL) &&
            getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_TIMING_SUMMARY") == NULL;
        checked = 1;
    }
    return enabled;
}

static uint32_t ds4_gpu_stream_expert_popcount(uint32_t mask) {
    return (uint32_t)__builtin_popcount(mask);
}

static int ds4_gpu_stream_expert_split_worthwhile(
        uint32_t resident_mask,
        uint32_t missing_mask) {
    if (resident_mask == 0 || missing_mask == 0) return 0;
    /*
     * The split path pays an extra command stage and a second routed-expert
     * bind.  It is worthwhile when several experts are missing and their SSD
     * reads can be hidden by resident expert work.  With one or two misses,
     * especially in large caches, a single unsplit routed pass is faster.
     */
    return ds4_gpu_stream_expert_popcount(missing_mask) >= 3u;
}

static void ds4_gpu_stream_expert_timing_note_selected(
        double sync_ms,
        double copy_ms,
        double bind_ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_selected_calls++;
    g_stream_expert_timing_selected_sync_ms += sync_ms;
    g_stream_expert_timing_selected_copy_ms += copy_ms;
    g_stream_expert_timing_selected_read_ms += sync_ms + copy_ms;
    g_stream_expert_timing_selected_bind_ms += bind_ms;
}

static void ds4_gpu_stream_expert_timing_note_split(
        uint32_t resident_mask,
        uint32_t missing_mask,
        double   resident_ms,
        double   missing_ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_split_layers++;
    g_stream_expert_timing_split_resident_experts +=
        ds4_gpu_stream_expert_popcount(resident_mask);
    g_stream_expert_timing_split_missing_experts +=
        ds4_gpu_stream_expert_popcount(missing_mask);
    g_stream_expert_timing_split_resident_ms += resident_ms;
    g_stream_expert_timing_split_missing_ms += missing_ms;
}

static void ds4_gpu_stream_expert_timing_note_split_missing_detail(
        double load_ms,
        double slot_ms,
        double prune_ms,
        double addr_ms,
        double wait_ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_split_missing_load_ms += load_ms;
    g_stream_expert_timing_split_missing_slot_ms += slot_ms;
    g_stream_expert_timing_split_missing_prune_ms += prune_ms;
    g_stream_expert_timing_split_missing_addr_ms += addr_ms;
    g_stream_expert_timing_split_missing_wait_ms += wait_ms;
}

static void ds4_gpu_stream_expert_timing_note_load_detail(
        double prepare_ms,
        double pread_ms,
        double modify_ms,
        double install_ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_load_calls++;
    g_stream_expert_timing_load_prepare_ms += prepare_ms;
    g_stream_expert_timing_load_pread_ms += pread_ms;
    g_stream_expert_timing_load_modify_ms += modify_ms;
    g_stream_expert_timing_load_install_ms += install_ms;
}

static void ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(double ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_prepare_batch_reuse_calls++;
    g_stream_expert_timing_prepare_batch_reuse_ms += ms;
}

static void ds4_gpu_stream_expert_timing_note_prepare_buffer(double ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_prepare_buffer_calls++;
    g_stream_expert_timing_prepare_buffer_ms += ms;
}

static void ds4_gpu_stream_expert_timing_note_prepare_task(
        uint32_t experts,
        double   ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_prepare_task_experts += experts;
    g_stream_expert_timing_prepare_task_ms += ms;
}

static void ds4_gpu_stream_expert_timing_note_reuse_scan(
        uint64_t entries,
        double   ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_reuse_scan_calls++;
    g_stream_expert_timing_reuse_scan_entries += entries;
    g_stream_expert_timing_reuse_scan_ms += ms;
}

static void ds4_gpu_stream_expert_timing_note_reuse_clear(double ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_reuse_clear_ms += ms;
}

static void ds4_gpu_stream_expert_timing_note_cache_class(
        uint32_t resident_mask,
        uint32_t missing_mask) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    const uint32_t resident = ds4_gpu_stream_expert_popcount(resident_mask);
    const uint32_t missing = ds4_gpu_stream_expert_popcount(missing_mask);
    if (missing == 0) {
        g_stream_expert_timing_cache_all_resident_layers++;
    } else if (resident == 0) {
        g_stream_expert_timing_cache_all_missing_layers++;
    } else {
        g_stream_expert_timing_cache_mixed_layers++;
    }
    g_stream_expert_timing_cache_resident_experts += resident;
    g_stream_expert_timing_cache_missing_experts += missing;
}

static int ds4_gpu_stream_expert_readahead_enabled(void) {
    return g_ssd_streaming_mode &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_READAHEAD") == NULL;
}

static void ds4_gpu_stream_expert_readahead_range(uint64_t offset, uint64_t len) {
    if (!ds4_gpu_stream_expert_readahead_enabled() || g_model_fd < 0 || len == 0) {
        return;
    }

#if defined(F_RDADVISE)
    const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
    const double t0 = timing ? ds4_gpu_now_ms() : 0.0;
    uint64_t pos = offset;
    uint64_t rem = len;
    while (rem > 0) {
        const uint64_t chunk64 =
            rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem;
        if (pos > (uint64_t)LLONG_MAX) break;

        struct radvisory ra;
        ra.ra_offset = (off_t)pos;
        ra.ra_count = (int)chunk64;
        (void)fcntl(g_model_fd, F_RDADVISE, &ra);

        pos += chunk64;
        rem -= chunk64;
    }
    if (timing) {
        g_stream_expert_timing_readahead_calls++;
        g_stream_expert_timing_readahead_bytes += len;
        g_stream_expert_timing_readahead_ms += ds4_gpu_now_ms() - t0;
    }
#else
    (void)offset;
    (void)len;
#endif
}

typedef struct {
    uint64_t offset;
    uint64_t len;
    uint8_t *dst;
    uint64_t read_bytes;
    double ms;
    int ok;
} ds4_gpu_stream_expert_pread_task;

typedef struct {
    int active;
    const void *model_map;
    uint64_t model_size;
    uint32_t layer;
    uint32_t n_total_expert;
    uint32_t n_selected;
    uint32_t missing_mask;
    uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint32_t n_loads;
    uint32_t n_tasks;
    uint64_t gate_expert_bytes;
    uint64_t down_expert_bytes;
    int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint64_t gate_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint64_t up_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint64_t down_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    __strong id<MTLBuffer> gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    __strong id<MTLBuffer> up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    __strong id<MTLBuffer> down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u];
    double start_ms;
    double prepare_ms;
} ds4_gpu_stream_expert_pending_load;

static ds4_gpu_stream_expert_pending_load g_stream_expert_pending_load;

typedef struct {
    int active;
    const void *model_map;
    uint64_t model_size;
    uint32_t layer;
    uint32_t n_total_expert;
    uint32_t n_selected;
    uint64_t gate_offset;
    uint64_t up_offset;
    uint64_t down_offset;
    uint64_t gate_expert_bytes;
    uint64_t down_expert_bytes;
    int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
} ds4_gpu_glm_stream_selected_prefetch;

static ds4_gpu_glm_stream_selected_prefetch g_glm_stream_selected_prefetch;

typedef struct {
    ds4_gpu_stream_expert_pread_task *tasks;
    uint32_t n_tasks;
    uint32_t worker_index;
    uint32_t n_workers;
} ds4_gpu_stream_expert_pread_worker_args;

static void ds4_gpu_stream_expert_cache_note_pread(
        uint32_t layer,
        uint64_t bytes,
        double   ms) {
    if (g_stream_expert_cache_pread_bytes > UINT64_MAX - bytes) {
        g_stream_expert_cache_pread_bytes = UINT64_MAX;
    } else {
        g_stream_expert_cache_pread_bytes += bytes;
    }
    g_stream_expert_cache_pread_ms += ms;
    if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) {
        if (g_stream_expert_cache_layer_pread_bytes[layer] > UINT64_MAX - bytes) {
            g_stream_expert_cache_layer_pread_bytes[layer] = UINT64_MAX;
        } else {
            g_stream_expert_cache_layer_pread_bytes[layer] += bytes;
        }
        g_stream_expert_cache_layer_pread_ms[layer] += ms;
    }
}

static uint32_t ds4_gpu_stream_expert_pread_thread_limit(void) {
    uint32_t threads = 9;
    const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_THREADS");
    if (env && env[0]) {
        char *end = NULL;
        unsigned long v = strtoul(env, &end, 10);
        if (end != env && *end == '\0') {
            threads = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v;
        }
    }
    if (threads == 0) threads = 1;
    if (threads > 18) threads = 18;
    return threads;
}

static uint32_t ds4_gpu_stream_expert_pread_thread_count(uint32_t n_tasks) {
    if (n_tasks <= 1) return n_tasks;
    uint32_t threads = ds4_gpu_stream_expert_pread_thread_limit();
    if (threads > n_tasks) threads = n_tasks;
    return threads;
}

static int ds4_gpu_stream_expert_pread_into(
        uint64_t  offset,
        uint64_t  len,
        uint8_t  *dst,
        uint64_t *read_bytes,
        double   *ms_out) {
    if (read_bytes) *read_bytes = 0;
    if (ms_out) *ms_out = 0.0;
    if (g_model_fd < 0 ||
        !dst ||
        len == 0 ||
        offset > (uint64_t)LLONG_MAX ||
        len > (uint64_t)LLONG_MAX - offset) {
        return 0;
    }

    const double t0 = ds4_gpu_now_ms();
    uint64_t pos = 0;
    int ok = 1;
    while (pos < len) {
        const uint64_t rem = len - pos;
        const size_t want = rem > (uint64_t)SSIZE_MAX ? (size_t)SSIZE_MAX : (size_t)rem;
        ssize_t nread;
        do {
            nread = pread(g_model_fd, dst + pos, want, (off_t)(offset + pos));
        } while (nread < 0 && errno == EINTR);
        if (nread <= 0) {
            ok = 0;
            break;
        }
        pos += (uint64_t)nread;
    }
    const double dt = ds4_gpu_now_ms() - t0;
    if (read_bytes) *read_bytes = pos;
    if (ms_out) *ms_out = dt;
    if (!ok || pos != len) {
        fprintf(stderr,
                "ds4: Metal streaming expert explicit pread failed offset=%.2f GiB len=%.2f MiB read=%.2f MiB\n",
                ds4_gpu_gib(offset),
                ds4_gpu_mib(len),
                ds4_gpu_mib(pos));
        return 0;
    }
    return 1;
}

static void *ds4_gpu_stream_expert_pread_worker(void *arg) {
    ds4_gpu_stream_expert_pread_worker_args *wa =
        (ds4_gpu_stream_expert_pread_worker_args *)arg;
    for (uint32_t i = wa->worker_index; i < wa->n_tasks; i += wa->n_workers) {
        ds4_gpu_stream_expert_pread_task *task = &wa->tasks[i];
        task->ok = ds4_gpu_stream_expert_pread_into(task->offset,
                                                    task->len,
                                                    task->dst,
                                                    &task->read_bytes,
                                                    &task->ms);
    }
    return NULL;
}

static pthread_mutex_t g_stream_expert_pread_pool_mutex = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t g_stream_expert_pread_pool_start_cond = PTHREAD_COND_INITIALIZER;
static pthread_cond_t g_stream_expert_pread_pool_done_cond = PTHREAD_COND_INITIALIZER;
static pthread_t g_stream_expert_pread_pool_threads[18];
static uint32_t g_stream_expert_pread_pool_thread_count;
static uint32_t g_stream_expert_pread_pool_active_workers;
static uint32_t g_stream_expert_pread_pool_remaining_workers;
static uint32_t g_stream_expert_pread_pool_n_tasks;
static uint32_t g_stream_expert_pread_pool_next_task;
static uint64_t g_stream_expert_pread_pool_generation;
static ds4_gpu_stream_expert_pread_task *g_stream_expert_pread_pool_tasks;
static int g_stream_expert_pread_pool_initialized;
static int g_stream_expert_pread_pool_stopping;

static int ds4_gpu_stream_expert_pread_pool_enabled(void) {
    const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_POOL");
    return !(env && strcmp(env, "0") == 0);
}

static void *ds4_gpu_stream_expert_pread_pool_worker(void *arg) {
    const uint32_t worker_index = (uint32_t)(uintptr_t)arg;
    uint64_t seen_generation = 0;

    for (;;) {
        pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
        while (!g_stream_expert_pread_pool_stopping &&
               g_stream_expert_pread_pool_generation == seen_generation) {
            pthread_cond_wait(&g_stream_expert_pread_pool_start_cond,
                              &g_stream_expert_pread_pool_mutex);
        }
        if (g_stream_expert_pread_pool_stopping) {
            pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
            break;
        }

        seen_generation = g_stream_expert_pread_pool_generation;
        if (worker_index >= g_stream_expert_pread_pool_active_workers) {
            pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
            continue;
        }

        for (;;) {
            const uint32_t task_index =
                g_stream_expert_pread_pool_next_task++;
            if (task_index >= g_stream_expert_pread_pool_n_tasks) break;

            ds4_gpu_stream_expert_pread_task *task =
                &g_stream_expert_pread_pool_tasks[task_index];
            pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);

            task->ok = ds4_gpu_stream_expert_pread_into(task->offset,
                                                        task->len,
                                                        task->dst,
                                                        &task->read_bytes,
                                                        &task->ms);

            pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
        }

        if (g_stream_expert_pread_pool_remaining_workers > 0 &&
            --g_stream_expert_pread_pool_remaining_workers == 0) {
            g_stream_expert_pread_pool_tasks = NULL;
            g_stream_expert_pread_pool_n_tasks = 0;
            g_stream_expert_pread_pool_active_workers = 0;
            pthread_cond_signal(&g_stream_expert_pread_pool_done_cond);
        }
        pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
    }

    return NULL;
}

static int ds4_gpu_stream_expert_pread_pool_init(uint32_t n_threads) {
    if (g_stream_expert_pread_pool_initialized) return 1;
    if (!ds4_gpu_stream_expert_pread_pool_enabled() || n_threads <= 1) return 0;
    if (n_threads > 18) n_threads = 18;

    pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
    g_stream_expert_pread_pool_thread_count = n_threads;
    g_stream_expert_pread_pool_stopping = 0;
    g_stream_expert_pread_pool_generation = 0;
    g_stream_expert_pread_pool_tasks = NULL;
    g_stream_expert_pread_pool_n_tasks = 0;
    g_stream_expert_pread_pool_next_task = 0;
    g_stream_expert_pread_pool_active_workers = 0;
    g_stream_expert_pread_pool_remaining_workers = 0;
    pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);

    uint32_t started = 0;
    for (uint32_t i = 0; i < n_threads; i++) {
        const int rc = pthread_create(&g_stream_expert_pread_pool_threads[i],
                                      NULL,
                                      ds4_gpu_stream_expert_pread_pool_worker,
                                      (void *)(uintptr_t)i);
        if (rc != 0) {
            fprintf(stderr,
                    "ds4: Metal streaming expert pread pool thread creation failed: %s\n",
                    strerror(rc));
            pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
            g_stream_expert_pread_pool_stopping = 1;
            g_stream_expert_pread_pool_generation++;
            pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond);
            pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
            for (uint32_t j = 0; j < started; j++) {
                (void)pthread_join(g_stream_expert_pread_pool_threads[j], NULL);
            }
            pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
            g_stream_expert_pread_pool_thread_count = 0;
            g_stream_expert_pread_pool_stopping = 0;
            pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
            return 0;
        }
        started++;
    }

    g_stream_expert_pread_pool_initialized = 1;
    return 1;
}

static int ds4_gpu_stream_expert_pread_pool_begin(
        ds4_gpu_stream_expert_pread_task *tasks,
        uint32_t n_tasks,
        uint32_t n_workers) {
    if (n_workers <= 1) return 0;
    const uint32_t limit = ds4_gpu_stream_expert_pread_thread_limit();
    if (!ds4_gpu_stream_expert_pread_pool_init(limit)) return 0;

    pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
    if (!g_stream_expert_pread_pool_initialized ||
        g_stream_expert_pread_pool_stopping ||
        g_stream_expert_pread_pool_thread_count == 0) {
        pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
        return 0;
    }
    if (n_workers > g_stream_expert_pread_pool_thread_count) {
        n_workers = g_stream_expert_pread_pool_thread_count;
    }
    if (n_workers == 0) {
        pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
        return 0;
    }

    if (g_stream_expert_pread_pool_remaining_workers != 0 ||
        g_stream_expert_pread_pool_tasks != NULL) {
        pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
        return 0;
    }

    g_stream_expert_pread_pool_tasks = tasks;
    g_stream_expert_pread_pool_n_tasks = n_tasks;
    g_stream_expert_pread_pool_next_task = 0;
    g_stream_expert_pread_pool_active_workers = n_workers;
    g_stream_expert_pread_pool_remaining_workers = n_workers;
    g_stream_expert_pread_pool_generation++;
    pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond);
    pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
    return 1;
}

static int ds4_gpu_stream_expert_pread_pool_wait(void) {
    if (!g_stream_expert_pread_pool_initialized) return 0;

    pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
    while (g_stream_expert_pread_pool_remaining_workers != 0) {
        pthread_cond_wait(&g_stream_expert_pread_pool_done_cond,
                          &g_stream_expert_pread_pool_mutex);
    }
    pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
    return 1;
}

static int ds4_gpu_stream_expert_pread_pool_dispatch(
        ds4_gpu_stream_expert_pread_task *tasks,
        uint32_t n_tasks,
        uint32_t n_workers) {
    if (!ds4_gpu_stream_expert_pread_pool_begin(tasks, n_tasks, n_workers)) {
        return 0;
    }
    return ds4_gpu_stream_expert_pread_pool_wait();
}

static void ds4_gpu_stream_expert_pread_pool_shutdown(void) {
    if (!g_stream_expert_pread_pool_initialized) return;

    pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
    g_stream_expert_pread_pool_stopping = 1;
    g_stream_expert_pread_pool_generation++;
    pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond);
    pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);

    const uint32_t n_threads = g_stream_expert_pread_pool_thread_count;
    for (uint32_t i = 0; i < n_threads; i++) {
        (void)pthread_join(g_stream_expert_pread_pool_threads[i], NULL);
    }

    pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
    g_stream_expert_pread_pool_thread_count = 0;
    g_stream_expert_pread_pool_active_workers = 0;
    g_stream_expert_pread_pool_remaining_workers = 0;
    g_stream_expert_pread_pool_n_tasks = 0;
    g_stream_expert_pread_pool_next_task = 0;
    g_stream_expert_pread_pool_tasks = NULL;
    g_stream_expert_pread_pool_initialized = 0;
    g_stream_expert_pread_pool_stopping = 0;
    pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
}

static int ds4_gpu_stream_expert_pread_tasks(
        ds4_gpu_stream_expert_pread_task *tasks,
        uint32_t n_tasks,
        uint64_t *total_bytes,
        double *wall_ms) {
    if (total_bytes) *total_bytes = 0;
    if (wall_ms) *wall_ms = 0.0;
    if (!tasks || n_tasks == 0) return 1;

    const uint32_t n_workers =
        ds4_gpu_stream_expert_pread_thread_count(n_tasks);
    const double t0 = ds4_gpu_now_ms();
    int ok = 1;
    if (n_workers <= 1) {
        ds4_gpu_stream_expert_pread_worker_args wa = {
            .tasks = tasks,
            .n_tasks = n_tasks,
            .worker_index = 0,
            .n_workers = 1,
        };
        (void)ds4_gpu_stream_expert_pread_worker(&wa);
    } else if (!ds4_gpu_stream_expert_pread_pool_dispatch(tasks,
                                                          n_tasks,
                                                          n_workers)) {
        pthread_t threads[18];
        ds4_gpu_stream_expert_pread_worker_args args[18];
        uint32_t started = 0;
        for (uint32_t i = 0; i < n_workers; i++) {
            args[i].tasks = tasks;
            args[i].n_tasks = n_tasks;
            args[i].worker_index = i;
            args[i].n_workers = n_workers;
            const int rc = pthread_create(&threads[i],
                                          NULL,
                                          ds4_gpu_stream_expert_pread_worker,
                                          &args[i]);
            if (rc != 0) {
                fprintf(stderr,
                        "ds4: Metal streaming expert pread thread creation failed: %s\n",
                        strerror(rc));
                ok = 0;
                break;
            }
            started++;
        }
        for (uint32_t i = 0; i < started; i++) {
            if (pthread_join(threads[i], NULL) != 0) ok = 0;
        }
    }
    const double dt = ds4_gpu_now_ms() - t0;

    uint64_t bytes = 0;
    for (uint32_t i = 0; i < n_tasks; i++) {
        if (!tasks[i].ok) ok = 0;
        if (bytes > UINT64_MAX - tasks[i].read_bytes) {
            bytes = UINT64_MAX;
        } else {
            bytes += tasks[i].read_bytes;
        }
    }
    if (total_bytes) *total_bytes = bytes;
    if (wall_ms) *wall_ms = dt;
    return ok;
}

static id<MTLBuffer> ds4_gpu_stream_expert_alloc_buffer(
        uint64_t  len,
        NSString *label) {
    if (!g_device ||
        len == 0 ||
        len > (uint64_t)NSUIntegerMax) {
        return nil;
    }

    id<MTLBuffer> buffer = [g_device newBufferWithLength:(NSUInteger)len
                                                 options:MTLResourceStorageModeShared];
    if (!buffer) {
        fprintf(stderr,
                "ds4: Metal streaming expert explicit buffer allocation failed (%.2f MiB)\n",
                ds4_gpu_mib(len));
        return nil;
    }
    buffer.label = label;
    g_stream_expert_cache_buffer_allocs++;
    return buffer;
}

static int ds4_gpu_stream_expert_combined_buffer_enabled(void) {
    return g_ssd_streaming_mode &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_COMBINED_BUFFER") == NULL;
}

static int ds4_gpu_stream_expert_slab_enabled(void) {
    return ds4_gpu_stream_expert_combined_buffer_enabled() &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_SLABS") == NULL;
}

/*
 * Large PRO caches otherwise create thousands of small shared Metal buffers.
 * Slabs keep the buffer object set small while locking pages only for slots
 * that actually hold a streamed expert.
 */
static uint64_t ds4_gpu_stream_expert_slab_target_bytes(void) {
    const uint64_t mib = 1024ull * 1024ull;
    uint64_t target = 4096ull * mib;
    const char *env = getenv("DS4_METAL_STREAMING_EXPERT_SLAB_MB");
    if (env && env[0]) {
        char *end = NULL;
        unsigned long long v = strtoull(env, &end, 10);
        if (end != env && *end == '\0' && v != 0) {
            target = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib;
        }
    }
    return target;
}

static id<MTLBuffer> ds4_gpu_stream_expert_alloc_slab_buffer(
        uint64_t  len,
        NSString *label) {
    if (!g_device ||
        len == 0 ||
        len > (uint64_t)NSUIntegerMax) {
        return nil;
    }

    id<MTLBuffer> buffer = [g_device newBufferWithLength:(NSUInteger)len
                                                 options:MTLResourceStorageModeShared];
    if (!buffer) {
        fprintf(stderr,
                "ds4: Metal streaming expert slab allocation failed (%.2f MiB)\n",
                ds4_gpu_mib(len));
        return nil;
    }
    buffer.label = label;
    g_stream_expert_cache_buffer_allocs++;
    return buffer;
}

static int ds4_gpu_stream_expert_slab_slot_range(
        uint32_t slot,
        uint32_t *slab_index,
        uint64_t *slot_base) {
    for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) {
        const uint32_t start = g_stream_expert_cache_slab_start_slot[i];
        const uint32_t count = g_stream_expert_cache_slab_slot_count[i];
        if (slot < start || slot >= start + count) continue;
        if (slab_index) *slab_index = i;
        if (slot_base) {
            *slot_base =
                (uint64_t)(slot - start) * g_stream_expert_cache_slab_slot_bytes;
        }
        return 1;
    }
    return 0;
}

static int ds4_gpu_stream_expert_slab_slot_for_buffer(
        id<MTLBuffer> buffer,
        NSUInteger    gate_inner,
        uint32_t     *slot_out) {
    if (!buffer || g_stream_expert_cache_slab_slot_bytes == 0 || !slot_out) {
        return 0;
    }
    for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) {
        if (g_stream_expert_cache_slabs[i] != buffer) continue;
        const uint64_t inner = (uint64_t)gate_inner;
        const uint64_t slot_bytes = g_stream_expert_cache_slab_slot_bytes;
        if (slot_bytes == 0 || inner % slot_bytes != 0) return 0;
        const uint64_t local_slot = inner / slot_bytes;
        if (local_slot >= g_stream_expert_cache_slab_slot_count[i]) return 0;
        const uint64_t slot =
            (uint64_t)g_stream_expert_cache_slab_start_slot[i] + local_slot;
        if (slot > UINT32_MAX) return 0;
        *slot_out = (uint32_t)slot;
        return 1;
    }
    return 0;
}

static void ds4_gpu_stream_expert_slab_push_free_slot(uint32_t slot) {
    if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
        g_stream_expert_cache_free_slot_count >=
            DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
        return;
    }
    g_stream_expert_cache_free_slots[g_stream_expert_cache_free_slot_count++] =
        slot;
}

static int ds4_gpu_stream_expert_slab_slot_buffers(
        uint32_t slot,
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes,
        __strong id<MTLBuffer> *gate_buf,
        __strong id<MTLBuffer> *up_buf,
        __strong id<MTLBuffer> *down_buf,
        NSUInteger *gate_inner,
        NSUInteger *up_inner,
        NSUInteger *down_inner) {
    uint32_t slab = UINT32_MAX;
    uint64_t base = 0;
    if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) ||
        slab >= g_stream_expert_cache_slab_count ||
        !g_stream_expert_cache_slabs[slab]) {
        return 0;
    }
    if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull ||
        base > UINT64_MAX - (gate_expert_bytes * 2ull + down_expert_bytes) ||
        base + gate_expert_bytes * 2ull + down_expert_bytes >
            (uint64_t)NSUIntegerMax) {
        return 0;
    }
    id<MTLBuffer> b = g_stream_expert_cache_slabs[slab];
    *gate_buf = b;
    *up_buf = b;
    *down_buf = b;
    *gate_inner = (NSUInteger)base;
    *up_inner = (NSUInteger)(base + gate_expert_bytes);
    *down_inner = (NSUInteger)(base + gate_expert_bytes * 2ull);
    return 1;
}

static int ds4_gpu_stream_expert_alloc_slab_slot(
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes,
        __strong id<MTLBuffer> *gate_buf,
        __strong id<MTLBuffer> *up_buf,
        __strong id<MTLBuffer> *down_buf,
        NSUInteger *gate_inner,
        NSUInteger *up_inner,
        NSUInteger *down_inner) {
    if (!ds4_gpu_stream_expert_slab_enabled() ||
        !gate_buf || !up_buf || !down_buf ||
        !gate_inner || !up_inner || !down_inner ||
        gate_expert_bytes == 0 || down_expert_bytes == 0 ||
        gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
        return 0;
    }

    uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
    if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0;
    const uint64_t page = (uint64_t)getpagesize();
    if (page != 0) {
        slot_bytes = round_up_u64(slot_bytes, page);
        if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0;
    }
    if (g_stream_expert_cache_slab_slot_bytes != 0 &&
        g_stream_expert_cache_slab_slot_bytes != slot_bytes) {
        return 0;
    }
    g_stream_expert_cache_slab_slot_bytes = slot_bytes;

    if (g_stream_expert_cache_free_slot_count != 0) {
        const uint32_t slot =
            g_stream_expert_cache_free_slots[--g_stream_expert_cache_free_slot_count];
        return ds4_gpu_stream_expert_slab_slot_buffers(slot,
                                                       gate_expert_bytes,
                                                       down_expert_bytes,
                                                       gate_buf,
                                                       up_buf,
                                                       down_buf,
                                                       gate_inner,
                                                       up_inner,
                                                       down_inner);
    }

    uint32_t slab = g_stream_expert_cache_slab_count;
    if (slab != 0 &&
        g_stream_expert_cache_slab_slots_used[slab - 1] <
            g_stream_expert_cache_slab_slot_count[slab - 1]) {
        slab--;
    } else {
        if (g_stream_expert_cache_slab_count >=
            DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS) {
            return 0;
        }
        const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
        if (budget != 0 && g_stream_expert_cache_slab_total_slots >= budget) {
            return 0;
        }
        uint64_t target = ds4_gpu_stream_expert_slab_target_bytes();
        uint64_t slots64 = target / slot_bytes;
        if (slots64 == 0) slots64 = 1;
        if (slots64 > UINT32_MAX) slots64 = UINT32_MAX;
        uint32_t slots = (uint32_t)slots64;
        if (budget != 0) {
            const uint32_t remaining =
                budget - g_stream_expert_cache_slab_total_slots;
            if (slots > remaining) slots = remaining;
        }
        if (slots == 0) return 0;
        id<MTLBuffer> slab_buffer = nil;
        while (slots != 0) {
            if ((uint64_t)slots <= UINT64_MAX / slot_bytes &&
                (uint64_t)slots * slot_bytes <= (uint64_t)NSUIntegerMax) {
                slab_buffer =
                    ds4_gpu_stream_expert_alloc_slab_buffer(
                            (uint64_t)slots * slot_bytes,
                            @"ds4_stream_expert_slab");
                if (slab_buffer) break;
            }
            slots /= 2u;
        }
        if (!slab_buffer || slots == 0) return 0;

        slab = g_stream_expert_cache_slab_count++;
        g_stream_expert_cache_slabs[slab] = slab_buffer;
        g_stream_expert_cache_slab_start_slot[slab] =
            g_stream_expert_cache_slab_total_slots;
        g_stream_expert_cache_slab_slot_count[slab] = slots;
        g_stream_expert_cache_slab_slots_used[slab] = 0;
        g_stream_expert_cache_slab_total_slots += slots;
    }

    const uint32_t local_slot = g_stream_expert_cache_slab_slots_used[slab]++;
    const uint32_t slot =
        g_stream_expert_cache_slab_start_slot[slab] + local_slot;
    return ds4_gpu_stream_expert_slab_slot_buffers(slot,
                                                   gate_expert_bytes,
                                                   down_expert_bytes,
                                                   gate_buf,
                                                   up_buf,
                                                   down_buf,
                                                   gate_inner,
                                                   up_inner,
                                                   down_inner);
}

static uint64_t ds4_gpu_stream_expert_buffer_object_count(
        id<MTLBuffer> gate,
        id<MTLBuffer> up,
        id<MTLBuffer> down) {
    if (!gate || !up || !down) return 0;
    if (gate == up && gate == down) return 1;
    if (gate == up || gate == down || up == down) return 2;
    return 3;
}

static int ds4_gpu_stream_expert_evict_dontneed_enabled(void) {
    return g_ssd_streaming_mode &&
           getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_EVICT_DONTNEED") != NULL &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EVICT_DONTNEED") == NULL;
}

static void ds4_gpu_stream_expert_evict_dontneed_range(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    offset,
        uint64_t    len) {
    if (!ds4_gpu_stream_expert_evict_dontneed_enabled() ||
        !model_map ||
        model_size == 0 ||
        offset > model_size ||
        len == 0 ||
        len > model_size - offset) {
        return;
    }

#if defined(POSIX_MADV_DONTNEED)
    const uint64_t page = (uint64_t)getpagesize();
    const uint64_t page_offset = offset & ~(page - 1);
    const uint64_t leading = offset - page_offset;
    if (len > UINT64_MAX - leading ||
        leading + len > UINT64_MAX - (page - 1)) {
        return;
    }
    uint64_t advise_bytes = round_up_u64(leading + len, page);
    if (advise_bytes > model_size - page_offset) {
        advise_bytes = model_size - page_offset;
    }
    if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) return;

    const uintptr_t base = (uintptr_t)model_map;
    if (page_offset > (uint64_t)(UINTPTR_MAX - base)) return;
    void *addr = (void *)(base + (uintptr_t)page_offset);
    const int rc = posix_madvise(addr, (size_t)advise_bytes, POSIX_MADV_DONTNEED);
    if (rc == 0) {
        if (g_stream_expert_cache_evict_advise_bytes > UINT64_MAX - advise_bytes) {
            g_stream_expert_cache_evict_advise_bytes = UINT64_MAX;
        } else {
            g_stream_expert_cache_evict_advise_bytes += advise_bytes;
        }
    } else if (getenv("DS4_METAL_STREAMING_EXPERT_EVICT_DONTNEED_PROFILE") != NULL) {
        fprintf(stderr,
                "ds4: Metal streaming expert evict DONTNEED failed offset=%.2f GiB len=%.2f MiB: %s\n",
                ds4_gpu_gib(offset),
                ds4_gpu_mib(len),
                strerror(rc));
    }
#else
    (void)model_map;
    (void)model_size;
    (void)offset;
    (void)len;
#endif
}

static int ds4_gpu_stream_expert_split_requested(void) {
    return g_ssd_streaming_mode;
}

static uint32_t ds4_gpu_stream_expert_split_min_decode_tokens(void) {
    return 4;
}

static uint32_t ds4_gpu_stream_expert_split_min_cached(void) {
    uint32_t min_cached = 1024;
    const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
    if (budget != 0 && budget < min_cached * 2u) {
        min_cached = budget / 2u;
    }
    return min_cached;
}

static int ds4_gpu_stream_expert_split_ready(void) {
    if (!ds4_gpu_stream_expert_split_requested()) return 0;
    if (g_stream_expert_cache_decode_tokens <
        ds4_gpu_stream_expert_split_min_decode_tokens()) {
        return 0;
    }
    return g_stream_expert_cache_entry_count >=
           ds4_gpu_stream_expert_split_min_cached();
}

static void ds4_gpu_stream_expert_cache_decay_route_hotness(void) {
    for (uint32_t layer = 0;
         layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
         layer++) {
        for (uint32_t expert = 0;
             expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
             expert++) {
            g_stream_expert_cache_route_hotness[layer][expert] >>= 1;
        }
    }
}

void ds4_gpu_stream_expert_cache_reset_route_hotness(void) {
    memset(g_stream_expert_cache_route_hotness,
           0,
           sizeof(g_stream_expert_cache_route_hotness));
    g_stream_expert_cache_hotness_decay_token =
        g_stream_expert_cache_decode_tokens;
}

static void ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(void) {
    if (g_stream_expert_cache_decode_tokens == 0) return;
    if (g_stream_expert_cache_hotness_decay_token == 0) {
        g_stream_expert_cache_hotness_decay_token =
            g_stream_expert_cache_decode_tokens;
        return;
    }
    while (g_stream_expert_cache_decode_tokens -
           g_stream_expert_cache_hotness_decay_token >=
           DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS) {
        ds4_gpu_stream_expert_cache_decay_route_hotness();
        g_stream_expert_cache_hotness_decay_token +=
            DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS;
    }
}

static void ds4_gpu_stream_expert_cache_note_route_hotness(
        uint32_t layer,
        uint32_t expert,
        uint32_t amount) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        amount == 0) {
        return;
    }
    uint32_t *hotness = &g_stream_expert_cache_route_hotness[layer][expert];
    if (*hotness > UINT32_MAX - amount) {
        *hotness = UINT32_MAX;
    } else {
        *hotness += amount;
    }
}

static void ds4_gpu_stream_expert_cache_note_selected_hotness(
        uint32_t       layer,
        const int32_t *selected_ids,
        uint32_t       n_selected) {
    if (!selected_ids ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_selected == 0) {
        return;
    }
    ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
    for (uint32_t i = 0; i < n_selected; i++) {
        if (selected_ids[i] < 0 ||
            selected_ids[i] >=
                (int32_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
            continue;
        }
        ds4_gpu_stream_expert_cache_note_route_hotness(
                layer,
                (uint32_t)selected_ids[i],
                1);
    }
}

static void ds4_gpu_stream_expert_cache_note_frequency_hotness(
        uint32_t        layer,
        const uint32_t *frequency,
        uint32_t        n_total_expert) {
    if (!frequency || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) {
        return;
    }
    if (n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
        n_total_expert = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
    }
    ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
    for (uint32_t expert = 0; expert < n_total_expert; expert++) {
        ds4_gpu_stream_expert_cache_note_route_hotness(layer,
                                                       expert,
                                                       frequency[expert]);
    }
}

static void ds4_gpu_stream_expert_cache_note_token(uint32_t layer_index) {
    if (!g_ssd_streaming_mode || layer_index != 0 ||
        g_stream_expert_cache_decode_tokens == UINT64_MAX) {
        return;
    }
    g_stream_expert_cache_decode_tokens++;
    ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
}

static void ds4_gpu_stream_expert_cache_note_decode_token(void) {
    if (!g_ssd_streaming_mode ||
        g_stream_expert_cache_decode_tokens == UINT64_MAX) {
        return;
    }
    g_stream_expert_cache_decode_tokens++;
    ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
}

static int ds4_gpu_stream_compact_addr_requested(void) {
    return g_ssd_streaming_mode &&
           getenv("DS4_METAL_ENABLE_STREAMING_COMPACT_ADDR") != NULL &&
           getenv("DS4_METAL_DISABLE_STREAMING_COMPACT_ADDR") == NULL &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL;
}

static int ds4_gpu_stream_expert_addr_table_requested(void) {
    return g_ssd_streaming_mode &&
           (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL ||
            getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL ||
            getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL ||
            g_stream_prefill_batch_selected_addr_building ||
            g_glm_stream_expert_addr_table_building ||
            (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL &&
             getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL) ||
            ds4_gpu_stream_expert_split_requested()) &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL;
}

static int ds4_gpu_stream_expert_addr_table_kernel_requested(void) {
    return g_ssd_streaming_mode &&
           (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL ||
            getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL ||
            getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL ||
            ds4_gpu_stream_expert_split_ready()) &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL;
}

static int ds4_gpu_stream_expert_masked_addr_requested(void) {
    return g_ssd_streaming_mode &&
           (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL ||
            ds4_gpu_stream_expert_split_ready()) &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_MASKED_ADDR") == NULL;
}

static int ds4_gpu_stream_expert_hit_validator_requested(void) {
    return g_ssd_streaming_mode &&
           getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL &&
           getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_HIT_VALIDATOR") == NULL;
}

static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_max(
        uint32_t n_total_expert) {
    const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX");
    if (env && env[0]) {
        char *end = NULL;
        const long v = strtol(env, &end, 10);
        if (end != env) {
            if (v <= 0) return 0;
            if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX;
            return (uint32_t)v;
        }
    }
    if (n_total_expert == 384) return 800u;
    if (n_total_expert == 256) return 760u;
    return 0;
}

static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_min(
        uint32_t n_total_expert) {
    const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN");
    if (env && env[0]) {
        char *end = NULL;
        const long v = strtol(env, &end, 10);
        if (end != env) {
            if (v <= 0) return 0;
            if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX;
            return (uint32_t)v;
        }
    }
    if (n_total_expert == 384 || n_total_expert == 256) return 2u;
    return 0;
}

static int ds4_gpu_stream_prefill_batch_selected_addr_enabled(
        uint32_t n_tokens,
        uint32_t n_total_expert,
        uint32_t n_expert,
        uint32_t gate_type,
        uint32_t down_type) {
    if (!g_ssd_streaming_mode ||
        n_tokens <= 1 ||
        n_total_expert == 0 ||
        n_expert != 6 ||
        gate_type != DS4_METAL_TENSOR_IQ2_XXS ||
        down_type != DS4_METAL_TENSOR_Q2_K ||
        g_quality_mode ||
        getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL ||
        getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL ||
        getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ||
        getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) {
        return 0;
    }
    /* All unique experts for one layer must fit simultaneously because the
     * address-table kernels consume them in one dispatch. Once the global
     * cache fills, preparation reuses entries owned by other layers. */
    if (ds4_gpu_stream_expert_cache_configured_count() < n_total_expert) {
        return 0;
    }
    if (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL) {
        return 1;
    }
    const uint32_t max_tokens =
        ds4_gpu_stream_prefill_batch_selected_addr_auto_max(n_total_expert);
    const uint32_t min_tokens =
        ds4_gpu_stream_prefill_batch_selected_addr_auto_min(n_total_expert);
    return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens;
}

static int ds4_gpu_glm_streaming_prefill_full_layer_active(void) {
    return g_glm_streaming_prefill_full_layer_runtime ||
           getenv("DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER") != NULL;
}

static int ds4_gpu_stream_full_expert_addr_table_requested(void) {
    return g_ssd_streaming_mode &&
           getenv("DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") != NULL &&
           getenv("DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") == NULL;
}

static uint64_t ds4_gpu_buffer_address(id<MTLBuffer> buffer, NSUInteger inner) {
    if (!buffer) return 0;
#if TARGET_OS_OSX
    if (@available(macOS 13.0, *)) {
        return (uint64_t)[buffer gpuAddress] + (uint64_t)inner;
    }
#endif
    return 0;
}

static int ds4_gpu_stream_compact_addr_ensure_buffers(uint32_t layer) {
    if (!g_device || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 0;

    const NSUInteger addr_bytes = 6u * sizeof(uint64_t);
    const NSUInteger ids_bytes = 6u * sizeof(int32_t);
    for (uint32_t i = 0; i < 4; i++) {
        id<MTLBuffer> current = nil;
        NSUInteger bytes = addr_bytes;
        NSString *label = @"ds4_stream_compact_gate_addresses";
        switch (i) {
        case 0:
            current = g_stream_compact_gate_addr_buffers[layer];
            label = @"ds4_stream_compact_gate_addresses";
            break;
        case 1:
            current = g_stream_compact_up_addr_buffers[layer];
            label = @"ds4_stream_compact_up_addresses";
            break;
        case 2:
            current = g_stream_compact_down_addr_buffers[layer];
            label = @"ds4_stream_compact_down_addresses";
            break;
        default:
            current = g_stream_compact_selected_buffers[layer];
            bytes = ids_bytes;
            label = @"ds4_stream_compact_selected_ids";
            break;
        }
        if (current) continue;
        id<MTLBuffer> b = [g_device newBufferWithLength:bytes
                                                options:MTLResourceStorageModeShared];
        if (!b) {
            fprintf(stderr, "ds4: Metal streaming compact address buffer allocation failed\n");
            return 0;
        }
        b.label = label;
        memset([b contents], 0, bytes);
        [b didModifyRange:NSMakeRange(0, bytes)];
        switch (i) {
        case 0:
            g_stream_compact_gate_addr_buffers[layer] = b;
            break;
        case 1:
            g_stream_compact_up_addr_buffers[layer] = b;
            break;
        case 2:
            g_stream_compact_down_addr_buffers[layer] = b;
            break;
        default:
            g_stream_compact_selected_buffers[layer] = b;
            break;
        }
    }
    return 1;
}

static int ds4_gpu_stream_compact_addr_prepare(
        uint32_t layer,
        ds4_gpu_stream_expert_cache_entry * const entries[6],
        uint32_t n_entries,
        id<MTLBuffer> *gate_addrs,
        id<MTLBuffer> *up_addrs,
        id<MTLBuffer> *down_addrs,
        id<MTLBuffer> *selected_ids) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        !entries || !gate_addrs || !up_addrs || !down_addrs || !selected_ids ||
        n_entries == 0 || n_entries > 6) {
        return 0;
    }
    if (!ds4_gpu_stream_compact_addr_ensure_buffers(layer)) return 0;

    uint64_t gate_values[6] = {0, 0, 0, 0, 0, 0};
    uint64_t up_values[6] = {0, 0, 0, 0, 0, 0};
    uint64_t down_values[6] = {0, 0, 0, 0, 0, 0};
    int32_t slot_ids[6] = {0, 1, 2, 3, 4, 5};

    for (uint32_t i = 0; i < n_entries; i++) {
        ds4_gpu_stream_expert_cache_entry *e = entries[i];
        if (!e || !e->gate_buffer || !e->up_buffer || !e->down_buffer) {
            return 0;
        }
        gate_values[i] = ds4_gpu_buffer_address(e->gate_buffer, e->gate_inner);
        up_values[i] = ds4_gpu_buffer_address(e->up_buffer, e->up_inner);
        down_values[i] = ds4_gpu_buffer_address(e->down_buffer, e->down_inner);
        if (gate_values[i] == 0 || up_values[i] == 0 || down_values[i] == 0) {
            fprintf(stderr, "ds4: Metal streaming compact address path requires GPU addresses\n");
            return 0;
        }
    }

    const NSUInteger addr_bytes = 6u * sizeof(uint64_t);
    const NSUInteger ids_bytes = 6u * sizeof(int32_t);
    id<MTLBuffer> gb = g_stream_compact_gate_addr_buffers[layer];
    id<MTLBuffer> ub = g_stream_compact_up_addr_buffers[layer];
    id<MTLBuffer> db = g_stream_compact_down_addr_buffers[layer];
    id<MTLBuffer> ib = g_stream_compact_selected_buffers[layer];
    memcpy([gb contents], gate_values, addr_bytes);
    memcpy([ub contents], up_values, addr_bytes);
    memcpy([db contents], down_values, addr_bytes);
    memcpy([ib contents], slot_ids, ids_bytes);
    [gb didModifyRange:NSMakeRange(0, addr_bytes)];
    [ub didModifyRange:NSMakeRange(0, addr_bytes)];
    [db didModifyRange:NSMakeRange(0, addr_bytes)];
    [ib didModifyRange:NSMakeRange(0, ids_bytes)];

    *gate_addrs = gb;
    *up_addrs = ub;
    *down_addrs = db;
    *selected_ids = ib;
    return 1;
}

static int ds4_gpu_stream_selected_ids_prepare(
        uint32_t       layer,
        const int32_t *selected_ids,
        uint32_t       n_selected,
        id<MTLBuffer> *selected_buf,
        NSUInteger    *selected_off) {
    if (!g_device ||
        !selected_ids ||
        !selected_buf ||
        !selected_off ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_selected == 0 ||
        n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED) {
        return 0;
    }

    const NSUInteger bytes =
        (NSUInteger)DS4_METAL_MAX_ROUTED_EXPERT_USED * sizeof(int32_t);
    id<MTLBuffer> b = g_stream_selected_id_buffers[layer];
    if (!b) {
        b = [g_device newBufferWithLength:bytes
                                  options:MTLResourceStorageModeShared];
        if (!b) {
            fprintf(stderr, "ds4: Metal streaming selected-id buffer allocation failed\n");
            return 0;
        }
        b.label = @"ds4_stream_selected_ids";
        g_stream_selected_id_buffers[layer] = b;
    }

    int32_t ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = {0};
    memcpy(ids, selected_ids, (size_t)n_selected * sizeof(ids[0]));
    memcpy([b contents], ids, bytes);
    [b didModifyRange:NSMakeRange(0, bytes)];

    *selected_buf = b;
    *selected_off = 0;
    return 1;
}

static int ds4_gpu_stream_expert_cache_ensure_addr_buffers(uint32_t layer) {
    if (!g_device ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) {
        return 0;
    }

    const NSUInteger bytes =
        (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t);
    id<MTLBuffer> buffers[3] = {
        g_stream_expert_cache_gate_addr_buffers[layer],
        g_stream_expert_cache_up_addr_buffers[layer],
        g_stream_expert_cache_down_addr_buffers[layer],
    };
    for (uint32_t i = 0; i < 3; i++) {
        if (buffers[i]) continue;
        id<MTLBuffer> b = [g_device newBufferWithLength:bytes
                                                options:MTLResourceStorageModeShared];
        if (!b) {
            fprintf(stderr, "ds4: Metal streaming expert address table allocation failed\n");
            return 0;
        }
        b.label =
            i == 0 ? @"ds4_stream_expert_gate_addresses" :
            (i == 1 ? @"ds4_stream_expert_up_addresses" :
                      @"ds4_stream_expert_down_addresses");
        memset([b contents], 0, bytes);
        [b didModifyRange:NSMakeRange(0, bytes)];
        if (i == 0) {
            g_stream_expert_cache_gate_addr_buffers[layer] = b;
        } else if (i == 1) {
            g_stream_expert_cache_up_addr_buffers[layer] = b;
        } else {
            g_stream_expert_cache_down_addr_buffers[layer] = b;
        }
    }
    return 1;
}

static void ds4_gpu_stream_expert_cache_zero_addr_slot(uint32_t layer, uint32_t expert) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
        return;
    }

    id<MTLBuffer> buffers[3] = {
        g_stream_expert_cache_gate_addr_buffers[layer],
        g_stream_expert_cache_up_addr_buffers[layer],
        g_stream_expert_cache_down_addr_buffers[layer],
    };
    const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t);
    for (uint32_t i = 0; i < 3; i++) {
        if (!buffers[i]) continue;
        uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off);
        *addr = 0;
        [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))];
    }
}

static int ds4_gpu_stream_expert_cache_set_addr_slot_raw(
        uint32_t      layer,
        uint32_t      expert,
        id<MTLBuffer> gate_buffer,
        NSUInteger    gate_inner,
        id<MTLBuffer> up_buffer,
        NSUInteger    up_inner,
        id<MTLBuffer> down_buffer,
        NSUInteger    down_inner) {
    if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0;

    const uint64_t values[3] = {
        ds4_gpu_buffer_address(gate_buffer, gate_inner),
        ds4_gpu_buffer_address(up_buffer, up_inner),
        ds4_gpu_buffer_address(down_buffer, down_inner),
    };
    if (values[0] == 0 || values[1] == 0 || values[2] == 0) {
        fprintf(stderr, "ds4: Metal streaming expert address table requires GPU addresses\n");
        return 0;
    }

    id<MTLBuffer> buffers[3] = {
        g_stream_expert_cache_gate_addr_buffers[layer],
        g_stream_expert_cache_up_addr_buffers[layer],
        g_stream_expert_cache_down_addr_buffers[layer],
    };
    const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t);
    for (uint32_t i = 0; i < 3; i++) {
        uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off);
        *addr = values[i];
        [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))];
    }
    return 1;
}

static int ds4_gpu_stream_expert_cache_set_addr_slot(
        uint32_t      layer,
        uint32_t      expert,
        id<MTLBuffer> gate_buffer,
        NSUInteger    gate_inner,
        id<MTLBuffer> up_buffer,
        NSUInteger    up_inner,
        id<MTLBuffer> down_buffer,
        NSUInteger    down_inner) {
    if (!ds4_gpu_stream_expert_addr_table_requested()) return 1;
    return ds4_gpu_stream_expert_cache_set_addr_slot_raw(layer,
                                                         expert,
                                                         gate_buffer,
                                                         gate_inner,
                                                         up_buffer,
                                                         up_inner,
                                                         down_buffer,
                                                         down_inner);
}

static int ds4_gpu_stream_expert_cache_addr_buffers(
        uint32_t       layer,
        id<MTLBuffer> *gate,
        id<MTLBuffer> *up,
        id<MTLBuffer> *down) {
    if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0;
    if (gate) *gate = g_stream_expert_cache_gate_addr_buffers[layer];
    if (up) *up = g_stream_expert_cache_up_addr_buffers[layer];
    if (down) *down = g_stream_expert_cache_down_addr_buffers[layer];
    return g_stream_expert_cache_gate_addr_buffers[layer] &&
           g_stream_expert_cache_up_addr_buffers[layer] &&
           g_stream_expert_cache_down_addr_buffers[layer];
}

static void ds4_gpu_stream_full_expert_addr_clear_layer(uint32_t layer) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return;
    ds4_gpu_stream_expert_cache_entry *e = &g_stream_full_expert_addr_entry[layer];
    e->gate_buffer = nil;
    e->up_buffer = nil;
    e->down_buffer = nil;
    e->model_map = NULL;
    e->model_size = 0;
    e->gate_abs_offset = 0;
    e->up_abs_offset = 0;
    e->down_abs_offset = 0;
    e->gate_expert_bytes = 0;
    e->down_expert_bytes = 0;
    e->logical_bytes = 0;
    e->last_used = 0;
    e->use_count = 0;
    e->gate_inner = 0;
    e->up_inner = 0;
    e->down_inner = 0;
    e->slab_slot = 0;
    e->valid = 0;
    e->slab_backed = 0;
}

static int ds4_gpu_stream_full_expert_addr_table_prepare(
        const void    *model_map,
        uint64_t       model_size,
        uint32_t       layer,
        uint32_t       n_total_expert,
        uint64_t       gate_abs_offset,
        uint64_t       up_abs_offset,
        uint64_t       down_abs_offset,
        uint64_t       gate_expert_bytes,
        uint64_t       down_expert_bytes,
        id<MTLBuffer> *gate_addrs,
        id<MTLBuffer> *up_addrs,
        id<MTLBuffer> *down_addrs,
        ds4_gpu_stream_expert_cache_entry **entry_out) {
    if (!ds4_gpu_stream_full_expert_addr_table_requested()) return 0;
    if (!model_map || model_size == 0 ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_total_expert == 0 ||
        n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        gate_expert_bytes == 0 ||
        down_expert_bytes == 0 ||
        n_total_expert > UINT64_MAX / gate_expert_bytes ||
        n_total_expert > UINT64_MAX / down_expert_bytes) {
        return 0;
    }

    const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes;
    if (gate_abs_offset > model_size ||
        up_abs_offset > model_size ||
        down_abs_offset > model_size ||
        gate_tensor_bytes > model_size - gate_abs_offset ||
        gate_tensor_bytes > model_size - up_abs_offset ||
        down_tensor_bytes > model_size - down_abs_offset) {
        return 0;
    }

    ds4_gpu_stream_expert_cache_entry *entry =
        &g_stream_full_expert_addr_entry[layer];
    if (!entry->valid ||
        entry->model_map != model_map ||
        entry->model_size != model_size ||
        entry->gate_abs_offset != gate_abs_offset ||
        entry->up_abs_offset != up_abs_offset ||
        entry->down_abs_offset != down_abs_offset ||
        entry->gate_expert_bytes != gate_expert_bytes ||
        entry->down_expert_bytes != down_expert_bytes ||
        !entry->gate_buffer ||
        !entry->up_buffer ||
        !entry->down_buffer) {
        ds4_gpu_stream_full_expert_addr_clear_layer(layer);

        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        uint64_t down_inner = 0;
        id<MTLBuffer> gate_buf =
            ds4_gpu_wrap_model_exact_range_owned(model_map,
                                                 model_size,
                                                 gate_abs_offset,
                                                 gate_tensor_bytes,
                                                 &gate_inner);
        id<MTLBuffer> up_buf =
            ds4_gpu_wrap_model_exact_range_owned(model_map,
                                                 model_size,
                                                 up_abs_offset,
                                                 gate_tensor_bytes,
                                                 &up_inner);
        id<MTLBuffer> down_buf =
            ds4_gpu_wrap_model_exact_range_owned(model_map,
                                                 model_size,
                                                 down_abs_offset,
                                                 down_tensor_bytes,
                                                 &down_inner);
        if (!gate_buf || !up_buf || !down_buf) return 0;

        entry->gate_buffer = gate_buf;
        entry->up_buffer = up_buf;
        entry->down_buffer = down_buf;
        entry->model_map = model_map;
        entry->model_size = model_size;
        entry->gate_abs_offset = gate_abs_offset;
        entry->up_abs_offset = up_abs_offset;
        entry->down_abs_offset = down_abs_offset;
        entry->gate_expert_bytes = gate_expert_bytes;
        entry->down_expert_bytes = down_expert_bytes;
        entry->logical_bytes = gate_tensor_bytes * 2ull + down_tensor_bytes;
        entry->gate_inner = (NSUInteger)gate_inner;
        entry->up_inner = (NSUInteger)up_inner;
        entry->down_inner = (NSUInteger)down_inner;
        entry->valid = 1;

        if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) {
            ds4_gpu_stream_full_expert_addr_clear_layer(layer);
            return 0;
        }

        const uint64_t gate_base = ds4_gpu_buffer_address(entry->gate_buffer,
                                                          entry->gate_inner);
        const uint64_t up_base = ds4_gpu_buffer_address(entry->up_buffer,
                                                        entry->up_inner);
        const uint64_t down_base = ds4_gpu_buffer_address(entry->down_buffer,
                                                          entry->down_inner);
        if (gate_base == 0 || up_base == 0 || down_base == 0) {
            fprintf(stderr, "ds4: Metal full streaming expert address table requires GPU addresses\n");
            ds4_gpu_stream_full_expert_addr_clear_layer(layer);
            return 0;
        }

        id<MTLBuffer> buffers[3] = {
            g_stream_expert_cache_gate_addr_buffers[layer],
            g_stream_expert_cache_up_addr_buffers[layer],
            g_stream_expert_cache_down_addr_buffers[layer],
        };
        uint64_t *gate_addr = (uint64_t *)[buffers[0] contents];
        uint64_t *up_addr = (uint64_t *)[buffers[1] contents];
        uint64_t *down_addr = (uint64_t *)[buffers[2] contents];
        for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) {
            if (expert < n_total_expert) {
                gate_addr[expert] = gate_base + (uint64_t)expert * gate_expert_bytes;
                up_addr[expert] = up_base + (uint64_t)expert * gate_expert_bytes;
                down_addr[expert] = down_base + (uint64_t)expert * down_expert_bytes;
            } else {
                gate_addr[expert] = 0;
                up_addr[expert] = 0;
                down_addr[expert] = 0;
            }
        }
        const NSUInteger bytes =
            (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t);
        for (uint32_t i = 0; i < 3; i++) {
            [buffers[i] didModifyRange:NSMakeRange(0, bytes)];
        }
    }

    if (gate_addrs) *gate_addrs = g_stream_expert_cache_gate_addr_buffers[layer];
    if (up_addrs) *up_addrs = g_stream_expert_cache_up_addr_buffers[layer];
    if (down_addrs) *down_addrs = g_stream_expert_cache_down_addr_buffers[layer];
    if (entry_out) *entry_out = entry;
    return entry->valid &&
           g_stream_expert_cache_gate_addr_buffers[layer] &&
           g_stream_expert_cache_up_addr_buffers[layer] &&
           g_stream_expert_cache_down_addr_buffers[layer];
}

static id<MTLBuffer> ds4_gpu_stream_expert_validate_status_buffer(void) {
    if (!g_device) return nil;
    if (g_stream_expert_validate_status_buffer) {
        return g_stream_expert_validate_status_buffer;
    }

    const NSUInteger bytes =
        (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t);
    id<MTLBuffer> b = [g_device newBufferWithLength:bytes
                                            options:MTLResourceStorageModeShared];
    if (!b) {
        fprintf(stderr, "ds4: Metal streaming expert validator allocation failed\n");
        return nil;
    }
    b.label = @"ds4_stream_expert_validate_status";
    memset([b contents], 0, bytes);
    [b didModifyRange:NSMakeRange(0, bytes)];
    g_stream_expert_validate_status_buffer = b;
    return b;
}

static int ds4_gpu_encode_stream_expert_cache_validate(
        id<MTLCommandBuffer> cb,
        const ds4_gpu_stream_expert_validate_args *args,
        id<MTLBuffer> selected,
        NSUInteger selected_off,
        id<MTLBuffer> gate_addrs,
        id<MTLBuffer> up_addrs,
        id<MTLBuffer> down_addrs,
        id<MTLBuffer> status) {
    if (!cb || !args || !selected || !gate_addrs || !up_addrs || !down_addrs ||
        !status || !g_moe_stream_expert_cache_validate_pipeline ||
        args->n_total_expert == 0 || args->n_total_expert > 384 ||
        args->n_expert == 0 || args->n_expert > 6) {
        return 0;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_moe_stream_expert_cache_validate_pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:selected offset:selected_off atIndex:1];
    [enc setBuffer:gate_addrs offset:0 atIndex:2];
    [enc setBuffer:up_addrs offset:0 atIndex:3];
    [enc setBuffer:down_addrs offset:0 atIndex:4];
    [enc setBuffer:status offset:0 atIndex:5];
    [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_stream_expert_cache_validate_selected(
        const ds4_gpu_tensor *selected,
        id<MTLBuffer> gate_addrs,
        id<MTLBuffer> up_addrs,
        id<MTLBuffer> down_addrs,
        uint32_t n_total_expert,
        uint32_t n_expert,
        int32_t selected_ids[6],
        uint32_t *all_cached,
        uint32_t *miss_mask,
        uint32_t *invalid_mask) {
    if (!selected || !selected_ids || !all_cached || !miss_mask || !invalid_mask) {
        return 0;
    }
    id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
    id<MTLBuffer> status = ds4_gpu_stream_expert_validate_status_buffer();
    if (!selectedbuf || !status) return 0;

    const NSUInteger status_bytes =
        (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t);
    memset([status contents], 0, status_bytes);
    [status didModifyRange:NSMakeRange(0, status_bytes)];

    ds4_gpu_stream_expert_validate_args args = {
        .n_total_expert = n_total_expert,
        .n_expert = n_expert,
    };

    const int had_batch = g_batch_cb != nil;
    int owned = 0;
    id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
    if (!cb) return 0;
    if (!ds4_gpu_encode_stream_expert_cache_validate(cb,
                                                     &args,
                                                     selectedbuf,
                                                     ds4_gpu_tensor_offset(selected),
                                                     gate_addrs,
                                                     up_addrs,
                                                     down_addrs,
                                                     status)) {
        return 0;
    }

    if (had_batch) {
        if (ds4_gpu_end_commands() == 0) return 0;
    } else if (!ds4_gpu_finish_command_buffer(cb, owned,
                                              "streaming expert cache validator")) {
        return 0;
    }

    const uint32_t *words = (const uint32_t *)[status contents];
    *all_cached = words[0];
    *miss_mask = words[1];
    *invalid_mask = words[2];
    for (uint32_t i = 0; i < 6; i++) {
        selected_ids[i] = (int32_t)words[4 + i];
    }

    if (had_batch && ds4_gpu_begin_commands() == 0) return 0;
    return 1;
}

static void ds4_gpu_stream_expert_cache_clear_entry_internal(
        uint32_t                                  layer,
        uint32_t                                  expert,
        int                                       count_eviction,
        int                                       recycle_slab_slot,
        ds4_gpu_stream_expert_reusable_buffers   *reuse) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
        return;
    }
    if (reuse) {
        reuse->gate_buffer = nil;
        reuse->up_buffer = nil;
        reuse->down_buffer = nil;
        reuse->gate_inner = 0;
        reuse->up_inner = 0;
        reuse->down_inner = 0;
    }
    ds4_gpu_stream_expert_cache_entry *e =
        &g_stream_expert_cache[layer][expert];
    if (!e->valid) return;
    if (ds4_gpu_stream_expert_cache_entry_inflight(e)) {
        return;
    }

    const uint64_t bytes = e->logical_bytes;
    ds4_gpu_stream_expert_evict_dontneed_range(e->model_map,
                                               e->model_size,
                                               e->gate_abs_offset,
                                               e->gate_expert_bytes);
    ds4_gpu_stream_expert_evict_dontneed_range(e->model_map,
                                               e->model_size,
                                               e->up_abs_offset,
                                               e->gate_expert_bytes);
    ds4_gpu_stream_expert_evict_dontneed_range(e->model_map,
                                               e->model_size,
                                               e->down_abs_offset,
                                               e->down_expert_bytes);
    ds4_gpu_stream_expert_cache_zero_addr_slot(layer, expert);
    if (reuse) {
        reuse->gate_buffer = e->gate_buffer;
        reuse->up_buffer = e->up_buffer;
        reuse->down_buffer = e->down_buffer;
        reuse->gate_inner = e->gate_inner;
        reuse->up_inner = e->up_inner;
        reuse->down_inner = e->down_inner;
    } else if (e->slab_backed && recycle_slab_slot) {
        ds4_gpu_stream_expert_slab_push_free_slot(e->slab_slot);
    }
    e->gate_buffer = nil;
    e->up_buffer = nil;
    e->down_buffer = nil;
    e->model_map = NULL;
    e->model_size = 0;
    e->gate_abs_offset = 0;
    e->up_abs_offset = 0;
    e->down_abs_offset = 0;
    e->gate_expert_bytes = 0;
    e->down_expert_bytes = 0;
    e->logical_bytes = 0;
    e->last_used = 0;
    e->use_count = 0;
    e->gate_inner = 0;
    e->up_inner = 0;
    e->down_inner = 0;
    e->inflight_seq = 0;
    e->slab_slot = 0;
    e->valid = 0;
    e->slab_backed = 0;

    if (g_stream_expert_cache_layer_count[layer] > 0) {
        g_stream_expert_cache_layer_count[layer]--;
    }
    if (g_stream_expert_cache_entry_count > 0) {
        g_stream_expert_cache_entry_count--;
    }
    if (g_stream_expert_cache_bytes >= bytes) {
        g_stream_expert_cache_bytes -= bytes;
    } else {
        g_stream_expert_cache_bytes = 0;
    }
    if (count_eviction) {
        g_stream_expert_cache_evictions++;
        g_stream_expert_cache_layer_evictions[layer]++;
    }
}

static void ds4_gpu_stream_expert_cache_clear_entry(
        uint32_t layer,
        uint32_t expert,
        int      count_eviction) {
    ds4_gpu_stream_expert_cache_clear_entry_internal(layer,
                                                     expert,
                                                     count_eviction,
                                                     1,
                                                     NULL);
}

static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) {
    ds4_gpu_stream_expert_pending_load_clear();
    g_stream_expert_cache_done_seq = g_stream_expert_cache_cb_seq;
    g_stream_expert_cache_batch_seq = 0;
    g_stream_expert_cache_owned_seq = 0;
    g_stream_expert_cache_pending_max_seq = 0;
    for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) {
        for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) {
            g_stream_expert_cache[layer][expert].inflight_seq = 0;
            ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
        }
        ds4_gpu_stream_full_expert_addr_clear_layer(layer);
        g_stream_expert_cache_layer_count[layer] = 0;
        id<MTLBuffer> buffers[3] = {
            g_stream_expert_cache_gate_addr_buffers[layer],
            g_stream_expert_cache_up_addr_buffers[layer],
            g_stream_expert_cache_down_addr_buffers[layer],
        };
        for (uint32_t i = 0; i < 3; i++) {
            if (!buffers[i]) continue;
            const NSUInteger bytes =
                (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t);
            memset([buffers[i] contents], 0, bytes);
            [buffers[i] didModifyRange:NSMakeRange(0, bytes)];
        }
    }
    g_stream_expert_cache_bytes = 0;
    g_stream_expert_cache_entry_count = 0;
    for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) {
        g_stream_expert_cache_slabs[i] = nil;
        g_stream_expert_cache_slab_start_slot[i] = 0;
        g_stream_expert_cache_slab_slot_count[i] = 0;
        g_stream_expert_cache_slab_slots_used[i] = 0;
    }
    g_stream_expert_cache_slab_count = 0;
    g_stream_expert_cache_slab_total_slots = 0;
    g_stream_expert_cache_free_slot_count = 0;
    g_stream_expert_cache_slab_slot_bytes = 0;
    if (reset_stats) {
        g_stream_expert_cache_hits = 0;
        g_stream_expert_cache_misses = 0;
        g_stream_expert_cache_evictions = 0;
        g_stream_expert_cache_wraps = 0;
        g_stream_expert_cache_clock = 0;
        g_stream_expert_cache_evict_advise_bytes = 0;
        g_stream_expert_cache_willneed_advise_bytes = 0;
        g_stream_expert_cache_pread_bytes = 0;
        g_stream_expert_cache_pread_ms = 0.0;
        g_stream_expert_cache_buffer_allocs = 0;
        g_stream_expert_cache_buffer_reuses = 0;
        g_stream_expert_cache_decode_tokens = 0;
        g_stream_expert_cache_hotness_decay_token = 0;
        memset(g_stream_expert_cache_route_hotness,
               0,
               sizeof(g_stream_expert_cache_route_hotness));
        g_stream_expert_timing_selected_calls = 0;
        g_stream_expert_timing_selected_read_ms = 0.0;
        g_stream_expert_timing_selected_sync_ms = 0.0;
        g_stream_expert_timing_selected_copy_ms = 0.0;
        g_stream_expert_timing_selected_bind_ms = 0.0;
        g_stream_expert_timing_split_layers = 0;
        g_stream_expert_timing_split_resident_experts = 0;
        g_stream_expert_timing_split_missing_experts = 0;
        g_stream_expert_timing_split_resident_ms = 0.0;
        g_stream_expert_timing_split_missing_ms = 0.0;
        g_stream_expert_timing_split_missing_load_ms = 0.0;
        g_stream_expert_timing_split_missing_slot_ms = 0.0;
        g_stream_expert_timing_split_missing_prune_ms = 0.0;
        g_stream_expert_timing_split_missing_addr_ms = 0.0;
        g_stream_expert_timing_split_missing_wait_ms = 0.0;
        g_stream_expert_timing_load_calls = 0;
        g_stream_expert_timing_load_prepare_ms = 0.0;
        g_stream_expert_timing_load_pread_ms = 0.0;
        g_stream_expert_timing_load_modify_ms = 0.0;
        g_stream_expert_timing_load_install_ms = 0.0;
        g_stream_expert_timing_prepare_batch_reuse_calls = 0;
        g_stream_expert_timing_prepare_batch_reuse_ms = 0.0;
        g_stream_expert_timing_prepare_buffer_calls = 0;
        g_stream_expert_timing_prepare_buffer_ms = 0.0;
        g_stream_expert_timing_prepare_task_experts = 0;
        g_stream_expert_timing_prepare_task_ms = 0.0;
        g_stream_expert_timing_reuse_scan_calls = 0;
        g_stream_expert_timing_reuse_scan_entries = 0;
        g_stream_expert_timing_reuse_scan_ms = 0.0;
        g_stream_expert_timing_reuse_clear_ms = 0.0;
        g_stream_expert_timing_readahead_calls = 0;
        g_stream_expert_timing_readahead_bytes = 0;
        g_stream_expert_timing_readahead_ms = 0.0;
        g_stream_expert_timing_cache_all_resident_layers = 0;
        g_stream_expert_timing_cache_all_missing_layers = 0;
        g_stream_expert_timing_cache_mixed_layers = 0;
        g_stream_expert_timing_cache_resident_experts = 0;
        g_stream_expert_timing_cache_missing_experts = 0;
        g_stream_expert_timing_last_report =
            (ds4_gpu_stream_expert_timing_snapshot){0};
        memset(g_stream_expert_cache_layer_hits,
               0,
               sizeof(g_stream_expert_cache_layer_hits));
        memset(g_stream_expert_cache_layer_misses,
               0,
               sizeof(g_stream_expert_cache_layer_misses));
        memset(g_stream_expert_cache_layer_evictions,
               0,
               sizeof(g_stream_expert_cache_layer_evictions));
        memset(g_stream_expert_cache_layer_pread_bytes,
               0,
               sizeof(g_stream_expert_cache_layer_pread_bytes));
        memset(g_stream_expert_cache_layer_pread_ms,
               0,
               sizeof(g_stream_expert_cache_layer_pread_ms));
    }
}

static int ds4_gpu_stream_expert_cache_is_protected(
        uint32_t expert,
        const int32_t *protect_ids,
        uint32_t n_protect) {
    if (!protect_ids) return 0;
    for (uint32_t i = 0; i < n_protect; i++) {
        if (protect_ids[i] >= 0 && (uint32_t)protect_ids[i] == expert) {
            return 1;
        }
    }
    return 0;
}

static void ds4_gpu_stream_expert_cache_prune_layer(
        uint32_t layer,
        uint32_t n_total_expert,
        uint32_t n_selected,
        const int32_t *protect_ids,
        uint32_t n_protect) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return;
    uint32_t cap = ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                             n_total_expert,
                                                             n_selected);
    if (cap == 0) return;

    /*
     * Route hotness counts selected experts even when they miss.  Hit-count
     * LFU penalizes experts that are repeatedly selected but evicted before a
     * second hit, which keeps too many decode layers in the mixed-cache path.
     */
    while (g_stream_expert_cache_layer_count[layer] > cap) {
        uint32_t victim = UINT32_MAX;
        uint32_t lowest_hotness = UINT32_MAX;
        uint64_t oldest = UINT64_MAX;
        for (uint32_t expert = 0; expert < n_total_expert; expert++) {
            ds4_gpu_stream_expert_cache_entry *e =
                &g_stream_expert_cache[layer][expert];
            if (!e->valid ||
                ds4_gpu_stream_expert_cache_entry_inflight(e) ||
                ds4_gpu_stream_expert_cache_is_protected(expert, protect_ids, n_protect)) {
                continue;
            }
            const uint32_t hotness =
                g_stream_expert_cache_route_hotness[layer][expert];
            if (hotness < lowest_hotness ||
                (hotness == lowest_hotness && e->last_used < oldest)) {
                lowest_hotness = hotness;
                oldest = e->last_used;
                victim = expert;
            }
        }
        if (victim == UINT32_MAX) break;
        ds4_gpu_stream_expert_cache_clear_entry(layer, victim, 1);
    }
}

static int ds4_gpu_stream_expert_cache_entry_protected(
        uint32_t layer,
        uint32_t expert,
        uint32_t protect_layer,
        const int32_t *protect_ids,
        uint32_t n_protect) {
    if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
        expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
        ds4_gpu_stream_expert_cache_entry_inflight(
                &g_stream_expert_cache[layer][expert])) {
        return 1;
    }
    return layer == protect_layer &&
           ds4_gpu_stream_expert_cache_is_protected(expert,
                                                    protect_ids,
                                                    n_protect);
}

static int ds4_gpu_stream_expert_cache_entry_reusable(
        const ds4_gpu_stream_expert_cache_entry *e,
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes) {
    return e &&
           e->valid &&
           e->gate_buffer &&
           e->up_buffer &&
           e->down_buffer &&
           e->gate_expert_bytes == gate_expert_bytes &&
           e->down_expert_bytes == down_expert_bytes;
}

static int ds4_gpu_stream_expert_cache_take_reusable(
        int                                     force_reuse,
        uint32_t                                protect_layer,
        const int32_t                          *protect_ids,
        uint32_t                                n_protect,
        uint64_t                                gate_expert_bytes,
        uint64_t                                down_expert_bytes,
        ds4_gpu_stream_expert_reusable_buffers *reuse) {
    if (!reuse) return 0;
    reuse->gate_buffer = nil;
    reuse->up_buffer = nil;
    reuse->down_buffer = nil;
    reuse->gate_inner = 0;
    reuse->up_inner = 0;
    reuse->down_inner = 0;

    const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
    if (budget == 0 ||
        (!force_reuse && g_stream_expert_cache_entry_count < budget)) {
        return 0;
    }

    int waited_inflight = 0;
retry:
    ;
    uint32_t victim_layer = UINT32_MAX;
    uint32_t victim_expert = UINT32_MAX;
    uint32_t lowest_hotness = UINT32_MAX;
    uint64_t oldest = UINT64_MAX;
    int skipped_inflight = 0;
    const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
    const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0;
    uint64_t scan_entries = 0;
    for (uint32_t layer = 0;
         layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
         layer++) {
        for (uint32_t expert = 0;
             expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
             expert++) {
            scan_entries++;
            ds4_gpu_stream_expert_cache_entry *e =
                &g_stream_expert_cache[layer][expert];
            if (!ds4_gpu_stream_expert_cache_entry_reusable(e,
                                                            gate_expert_bytes,
                                                            down_expert_bytes)) {
                continue;
            }
            if (ds4_gpu_stream_expert_cache_entry_inflight(e)) {
                skipped_inflight = 1;
                continue;
            }
            if (ds4_gpu_stream_expert_cache_entry_protected(layer,
                                                            expert,
                                                            protect_layer,
                                                            protect_ids,
                                                            n_protect)) {
                continue;
            }
            const uint32_t hotness =
                g_stream_expert_cache_route_hotness[layer][expert];
            if (hotness < lowest_hotness ||
                (hotness == lowest_hotness && e->last_used < oldest)) {
                lowest_hotness = hotness;
                oldest = e->last_used;
                victim_layer = layer;
                victim_expert = expert;
            }
        }
    }
    if (timing) {
        ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries,
                                                     ds4_gpu_now_ms() - scan_t0);
    }

    if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) {
        if (skipped_inflight && !waited_inflight &&
            !ds4_gpu_stream_expert_cache_on_service_thread()) {
            waited_inflight = 1;
            if (!ds4_gpu_stream_expert_cache_wait_inflight(
                    "streaming expert cache reuse")) {
                return 0;
            }
            goto retry;
        }
        return 0;
    }
    const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0;
    ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer,
                                                     victim_expert,
                                                     1,
                                                     1,
                                                     reuse);
    if (timing) {
        ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() -
                                                      clear_t0);
    }
    if (!reuse->gate_buffer || !reuse->up_buffer || !reuse->down_buffer) {
        reuse->gate_buffer = nil;
        reuse->up_buffer = nil;
        reuse->down_buffer = nil;
        reuse->gate_inner = 0;
        reuse->up_inner = 0;
        reuse->down_inner = 0;
        return 0;
    }
    g_stream_expert_cache_buffer_reuses +=
        ds4_gpu_stream_expert_buffer_object_count(reuse->gate_buffer,
                                                  reuse->up_buffer,
                                                  reuse->down_buffer);
    return 1;
}

static uint32_t ds4_gpu_stream_expert_cache_take_reusable_batch(
        uint32_t                                n_needed,
        uint32_t                                protect_layer,
        const int32_t                          *protect_ids,
        uint32_t                                n_protect,
        uint64_t                                gate_expert_bytes,
        uint64_t                                down_expert_bytes,
        ds4_gpu_stream_expert_reusable_buffers *reuses) {
    if (!reuses || n_needed == 0) return 0;
    if (n_needed > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
        n_needed = DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED;
    }
    for (uint32_t i = 0; i < n_needed; i++) {
        reuses[i].gate_buffer = nil;
        reuses[i].up_buffer = nil;
        reuses[i].down_buffer = nil;
        reuses[i].gate_inner = 0;
        reuses[i].up_inner = 0;
        reuses[i].down_inner = 0;
    }

    const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
    if (budget == 0 || g_stream_expert_cache_entry_count < budget) {
        return 0;
    }

    int waited_inflight = 0;
retry:
    ;
    uint32_t victim_layers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint32_t victim_experts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint32_t victim_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint64_t victim_last_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint32_t victim_count = 0;
    int skipped_inflight = 0;
    const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
    const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0;
    uint64_t scan_entries = 0;
    for (uint32_t i = 0; i < n_needed; i++) {
        victim_layers[i] = UINT32_MAX;
        victim_experts[i] = UINT32_MAX;
        victim_hotness[i] = UINT32_MAX;
        victim_last_used[i] = UINT64_MAX;
    }

    for (uint32_t layer = 0;
         layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
         layer++) {
        for (uint32_t expert = 0;
             expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
             expert++) {
            scan_entries++;
            ds4_gpu_stream_expert_cache_entry *e =
                &g_stream_expert_cache[layer][expert];
            if (!ds4_gpu_stream_expert_cache_entry_reusable(e,
                                                            gate_expert_bytes,
                                                            down_expert_bytes)) {
                continue;
            }
            if (ds4_gpu_stream_expert_cache_entry_inflight(e)) {
                skipped_inflight = 1;
                continue;
            }
            if (ds4_gpu_stream_expert_cache_entry_protected(layer,
                                                            expert,
                                                            protect_layer,
                                                            protect_ids,
                                                            n_protect)) {
                continue;
            }

            const uint32_t hotness =
                g_stream_expert_cache_route_hotness[layer][expert];
            const uint64_t last_used = e->last_used;
            if (victim_count < n_needed) {
                victim_layers[victim_count] = layer;
                victim_experts[victim_count] = expert;
                victim_hotness[victim_count] = hotness;
                victim_last_used[victim_count] = last_used;
                victim_count++;
                continue;
            }

            uint32_t worst = 0;
            for (uint32_t i = 1; i < victim_count; i++) {
                if (victim_hotness[i] > victim_hotness[worst] ||
                    (victim_hotness[i] == victim_hotness[worst] &&
                     victim_last_used[i] > victim_last_used[worst])) {
                    worst = i;
                }
            }
            if (hotness < victim_hotness[worst] ||
                (hotness == victim_hotness[worst] &&
                 last_used < victim_last_used[worst])) {
                victim_layers[worst] = layer;
                victim_experts[worst] = expert;
                victim_hotness[worst] = hotness;
                victim_last_used[worst] = last_used;
            }
        }
    }
    if (timing) {
        ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries,
                                                     ds4_gpu_now_ms() - scan_t0);
    }

    if (victim_count == 0) {
        if (skipped_inflight && !waited_inflight) {
            waited_inflight = 1;
            if (!ds4_gpu_stream_expert_cache_wait_inflight(
                    "streaming expert cache batch reuse")) {
                return 0;
            }
            goto retry;
        }
        return 0;
    }

    uint32_t reuse_count = 0;
    const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0;
    for (uint32_t i = 0; i < victim_count; i++) {
        if (victim_layers[i] == UINT32_MAX ||
            victim_experts[i] == UINT32_MAX) {
            continue;
        }
        ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layers[i],
                                                         victim_experts[i],
                                                         1,
                                                         1,
                                                         &reuses[reuse_count]);
        if (!reuses[reuse_count].gate_buffer ||
            !reuses[reuse_count].up_buffer ||
            !reuses[reuse_count].down_buffer) {
            reuses[reuse_count].gate_buffer = nil;
            reuses[reuse_count].up_buffer = nil;
            reuses[reuse_count].down_buffer = nil;
            reuses[reuse_count].gate_inner = 0;
            reuses[reuse_count].up_inner = 0;
            reuses[reuse_count].down_inner = 0;
            continue;
        }
        g_stream_expert_cache_buffer_reuses +=
            ds4_gpu_stream_expert_buffer_object_count(
                    reuses[reuse_count].gate_buffer,
                    reuses[reuse_count].up_buffer,
                    reuses[reuse_count].down_buffer);
        reuse_count++;
    }
    if (timing) {
        ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() -
                                                      clear_t0);
    }
    return reuse_count;
}

static int ds4_gpu_stream_expert_batch_reuse_enabled(
        uint64_t gate_expert_bytes,
        uint64_t down_expert_bytes) {
    if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
        return 0;
    }
    const uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
    /*
     * One global victim scan per selected miss is measurable for GLM Q2-size
     * slots, but batching larger Q4-size slots regressed short decode on M5.
     * Keep larger slots on the older single-victim path until profiling says
     * otherwise.
     */
    return slot_bytes <= 16ull * 1024ull * 1024ull;
}

static int ds4_gpu_stream_expert_cache_prepare_load_buffers(
        uint32_t       layer,
        uint32_t       expert,
        uint32_t       protect_layer,
        const int32_t *protect_ids,
        uint32_t       n_protect,
        uint64_t       gate_expert_bytes,
        uint64_t       down_expert_bytes,
        int            force_reuse,
        __strong id<MTLBuffer> *gate_buf,
        __strong id<MTLBuffer> *up_buf,
        __strong id<MTLBuffer> *down_buf,
        NSUInteger    *gate_inner,
        NSUInteger    *up_inner,
        NSUInteger    *down_inner) {
    if (!gate_buf || !up_buf || !down_buf ||
        !gate_inner || !up_inner || !down_inner ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
        return 0;
    }

    *gate_buf = nil;
    *up_buf = nil;
    *down_buf = nil;
    *gate_inner = 0;
    *up_inner = 0;
    *down_inner = 0;

    ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 };
    ds4_gpu_stream_expert_cache_entry *e =
        &g_stream_expert_cache[layer][expert];
    if (e->valid && ds4_gpu_stream_expert_cache_entry_inflight(e)) {
        if (ds4_gpu_stream_expert_cache_on_service_thread()) return 0;
        if (!ds4_gpu_stream_expert_cache_wait_inflight(
                "streaming expert cache replacement")) {
            return 0;
        }
        e = &g_stream_expert_cache[layer][expert];
    }
    if (ds4_gpu_stream_expert_cache_entry_reusable(e,
                                                   gate_expert_bytes,
                                                   down_expert_bytes)) {
        ds4_gpu_stream_expert_cache_clear_entry_internal(layer,
                                                         expert,
                                                         0,
                                                         1,
                                                         &reuse);
        if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) {
            g_stream_expert_cache_buffer_reuses +=
                ds4_gpu_stream_expert_buffer_object_count(reuse.gate_buffer,
                                                          reuse.up_buffer,
                                                          reuse.down_buffer);
        }
    } else if (e->valid) {
        ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
    }

    if (!reuse.gate_buffer || !reuse.up_buffer || !reuse.down_buffer) {
        if (!ds4_gpu_stream_expert_cache_take_reusable(force_reuse,
                                                       protect_layer,
                                                       protect_ids,
                                                       n_protect,
                                                       gate_expert_bytes,
                                                       down_expert_bytes,
                                                       &reuse)) {
            reuse.gate_buffer = nil;
            reuse.up_buffer = nil;
            reuse.down_buffer = nil;
        }
    }

    if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) {
        *gate_buf = reuse.gate_buffer;
        *up_buf = reuse.up_buffer;
        *down_buf = reuse.down_buffer;
        *gate_inner = reuse.gate_inner;
        *up_inner = reuse.up_inner;
        *down_inner = reuse.down_inner;
        return 1;
    }

    if (ds4_gpu_stream_expert_combined_buffer_enabled()) {
        if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes ||
            gate_expert_bytes * 2ull > UINT64_MAX - down_expert_bytes ||
            gate_expert_bytes > (uint64_t)NSUIntegerMax ||
            gate_expert_bytes * 2ull > (uint64_t)NSUIntegerMax ||
            gate_expert_bytes * 2ull + down_expert_bytes >
                (uint64_t)NSUIntegerMax) {
            return 0;
        }
        if (ds4_gpu_stream_expert_alloc_slab_slot(gate_expert_bytes,
                                                  down_expert_bytes,
                                                  gate_buf,
                                                  up_buf,
                                                  down_buf,
                                                  gate_inner,
                                                  up_inner,
                                                  down_inner)) {
            return 1;
        }
        const uint64_t up_off = gate_expert_bytes;
        const uint64_t down_off = gate_expert_bytes * 2ull;
        const uint64_t combined_bytes = down_off + down_expert_bytes;
        id<MTLBuffer> combined =
            ds4_gpu_stream_expert_alloc_buffer(combined_bytes,
                                               @"ds4_stream_expert_combined");
        if (!combined) return 0;
        *gate_buf = combined;
        *up_buf = combined;
        *down_buf = combined;
        *gate_inner = 0;
        *up_inner = (NSUInteger)up_off;
        *down_inner = (NSUInteger)down_off;
        return 1;
    }

    *gate_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes,
                                                   @"ds4_stream_expert_gate");
    *up_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes,
                                                 @"ds4_stream_expert_up");
    *down_buf = ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes,
                                                   @"ds4_stream_expert_down");
    return *gate_buf && *up_buf && *down_buf;
}

static void ds4_gpu_stream_expert_cache_prune_global(
        uint32_t protect_layer,
        const int32_t *protect_ids,
        uint32_t n_protect) {
    const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
    if (budget == 0 || g_stream_expert_cache_entry_count <= budget) return;

    while (g_stream_expert_cache_entry_count > budget) {
        uint32_t victim_layer = UINT32_MAX;
        uint32_t victim_expert = UINT32_MAX;
        uint32_t lowest_hotness = UINT32_MAX;
        uint64_t oldest = UINT64_MAX;
        for (uint32_t layer = 0;
             layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
             layer++) {
            for (uint32_t expert = 0;
                 expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
                 expert++) {
                ds4_gpu_stream_expert_cache_entry *e =
                    &g_stream_expert_cache[layer][expert];
                if (!e->valid ||
                    ds4_gpu_stream_expert_cache_entry_protected(layer,
                                                                expert,
                                                                protect_layer,
                                                                protect_ids,
                                                                n_protect)) {
                    continue;
                }
                const uint32_t hotness =
                    g_stream_expert_cache_route_hotness[layer][expert];
                if (hotness < lowest_hotness ||
                    (hotness == lowest_hotness && e->last_used < oldest)) {
                    lowest_hotness = hotness;
                    oldest = e->last_used;
                    victim_layer = layer;
                    victim_expert = expert;
                }
            }
        }
        if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) break;
        ds4_gpu_stream_expert_cache_clear_entry(victim_layer, victim_expert, 1);
    }
}

static int ds4_gpu_stream_expert_cache_entry_matches(
        const ds4_gpu_stream_expert_cache_entry *e,
        const void *model_map,
        uint64_t    model_size,
        uint64_t    gate_abs_offset,
        uint64_t    up_abs_offset,
        uint64_t    down_abs_offset,
        uint64_t    gate_expert_bytes,
        uint64_t    down_expert_bytes) {
    return e &&
           e->valid &&
           e->model_map == model_map &&
           e->model_size == model_size &&
           e->gate_abs_offset == gate_abs_offset &&
           e->up_abs_offset == up_abs_offset &&
           e->down_abs_offset == down_abs_offset &&
           e->gate_expert_bytes == gate_expert_bytes &&
           e->down_expert_bytes == down_expert_bytes &&
           e->gate_buffer && e->up_buffer && e->down_buffer;
}

static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_peek(
        const void *model_map,
        uint64_t    model_size,
        uint32_t    layer,
        uint32_t    expert,
        uint32_t    n_total_expert,
        uint32_t    n_selected,
        uint64_t    gate_abs_offset,
        uint64_t    up_abs_offset,
        uint64_t    down_abs_offset,
        uint64_t    gate_expert_bytes,
        uint64_t    down_expert_bytes) {
    if (!g_ssd_streaming_mode ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        expert >= n_total_expert ||
        !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes) ||
        ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                  n_total_expert,
                                                  n_selected) == 0) {
        return NULL;
    }

    ds4_gpu_stream_expert_cache_entry *e =
        &g_stream_expert_cache[layer][expert];
    if (!ds4_gpu_stream_expert_cache_entry_matches(e,
                                                   model_map,
                                                   model_size,
                                                   gate_abs_offset,
                                                   up_abs_offset,
                                                   down_abs_offset,
                                                   gate_expert_bytes,
                                                   down_expert_bytes)) {
        return NULL;
    }

    e->last_used = ++g_stream_expert_cache_clock;
    e->use_count++;
    g_stream_expert_cache_hits++;
    g_stream_expert_cache_layer_hits[layer]++;
    return e;
}

static ds4_gpu_stream_expert_cache_entry *
ds4_gpu_stream_expert_cache_install_loaded(
        const void    *model_map,
        uint64_t       model_size,
        uint32_t       layer,
        uint32_t       expert,
        uint64_t       gate_abs_offset,
        uint64_t       up_abs_offset,
        uint64_t       down_abs_offset,
        uint64_t       gate_expert_bytes,
        uint64_t       down_expert_bytes,
        id<MTLBuffer>  gate_buf,
        id<MTLBuffer>  up_buf,
        id<MTLBuffer>  down_buf,
        NSUInteger     gate_inner,
        NSUInteger     up_inner,
        NSUInteger     down_inner) {
    if (!gate_buf || !up_buf || !down_buf ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
        return NULL;
    }
    if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
        fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n");
        return NULL;
    }
    const uint64_t logical_bytes = gate_expert_bytes * 2ull + down_expert_bytes;

    ds4_gpu_stream_expert_cache_entry *e =
        &g_stream_expert_cache[layer][expert];
    if (e->valid) {
        if (ds4_gpu_stream_expert_cache_entry_inflight(e) &&
            ds4_gpu_stream_expert_cache_on_service_thread()) {
            return NULL;
        }
        if (ds4_gpu_stream_expert_cache_entry_inflight(e) &&
            !ds4_gpu_stream_expert_cache_wait_inflight(
                    "streaming expert cache install")) {
            return NULL;
        }
        if (ds4_gpu_stream_expert_cache_entry_inflight(e)) return NULL;
        ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
        if (e->valid) return NULL;
    }

    if (!ds4_gpu_stream_expert_cache_set_addr_slot(layer,
                                                   expert,
                                                   gate_buf,
                                                   gate_inner,
                                                   up_buf,
                                                   up_inner,
                                                   down_buf,
                                                   down_inner)) {
        return NULL;
    }

    e->gate_buffer = gate_buf;
    e->up_buffer = up_buf;
    e->down_buffer = down_buf;
    e->model_map = model_map;
    e->model_size = model_size;
    e->gate_abs_offset = gate_abs_offset;
    e->up_abs_offset = up_abs_offset;
    e->down_abs_offset = down_abs_offset;
    e->gate_expert_bytes = gate_expert_bytes;
    e->down_expert_bytes = down_expert_bytes;
    e->logical_bytes = logical_bytes;
    e->last_used = ++g_stream_expert_cache_clock;
    e->use_count = 1;
    e->gate_inner = gate_inner;
    e->up_inner = up_inner;
    e->down_inner = down_inner;
    e->inflight_seq = 0;
    uint32_t slab_slot = 0;
    if (gate_buf == up_buf &&
        gate_buf == down_buf &&
        ds4_gpu_stream_expert_slab_slot_for_buffer(gate_buf,
                                                   gate_inner,
                                                   &slab_slot)) {
        e->slab_backed = 1;
        e->slab_slot = slab_slot;
    } else {
        e->slab_backed = 0;
        e->slab_slot = 0;
    }
    e->valid = 1;
    g_stream_expert_cache_layer_count[layer]++;
    if (g_stream_expert_cache_entry_count < UINT32_MAX) {
        g_stream_expert_cache_entry_count++;
    }
    if (g_stream_expert_cache_bytes > UINT64_MAX - logical_bytes) {
        g_stream_expert_cache_bytes = UINT64_MAX;
    } else {
        g_stream_expert_cache_bytes += logical_bytes;
    }
    g_stream_expert_cache_misses++;
    g_stream_expert_cache_layer_misses[layer]++;
    g_stream_expert_cache_wraps += 3;
    return e;
}

static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get_protected(
        const void *model_map,
        uint64_t    model_size,
        uint32_t    layer,
        uint32_t    expert,
        uint32_t    n_total_expert,
        uint32_t    n_selected,
        uint64_t    gate_abs_offset,
        uint64_t    up_abs_offset,
        uint64_t    down_abs_offset,
        uint64_t    gate_expert_bytes,
        uint64_t    down_expert_bytes,
        const int32_t *protect_ids,
        uint32_t    n_protect) {
    if (!g_ssd_streaming_mode ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        expert >= n_total_expert ||
        !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes) ||
        ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                  n_total_expert,
                                                  n_selected) == 0) {
        return NULL;
    }

    ds4_gpu_stream_expert_cache_entry *e =
        &g_stream_expert_cache[layer][expert];
    if (ds4_gpu_stream_expert_cache_entry_matches(e,
                                                  model_map,
                                                  model_size,
                                                  gate_abs_offset,
                                                  up_abs_offset,
                                                  down_abs_offset,
                                                  gate_expert_bytes,
                                                  down_expert_bytes)) {
        e->last_used = ++g_stream_expert_cache_clock;
        e->use_count++;
        g_stream_expert_cache_hits++;
        g_stream_expert_cache_layer_hits[layer]++;
        return e;
    }

    ds4_gpu_stream_expert_readahead_range(gate_abs_offset, gate_expert_bytes);
    ds4_gpu_stream_expert_readahead_range(up_abs_offset, gate_expert_bytes);
    ds4_gpu_stream_expert_readahead_range(down_abs_offset, down_expert_bytes);

    id<MTLBuffer> gate_buf = nil;
    id<MTLBuffer> up_buf = nil;
    id<MTLBuffer> down_buf = nil;
    NSUInteger gate_inner = 0;
    NSUInteger up_inner = 0;
    NSUInteger down_inner = 0;
    const int32_t protect_one = (int32_t)expert;
    if (!protect_ids || n_protect == 0) {
        protect_ids = &protect_one;
        n_protect = 1;
    }
    const uint32_t cache_budget =
        ds4_gpu_stream_expert_cache_configured_budget();
    const int force_reuse =
        cache_budget != 0 && g_stream_expert_cache_entry_count >= cache_budget;
    if (!ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
                                                          expert,
                                                          layer,
                                                          protect_ids,
                                                          n_protect,
                                                          gate_expert_bytes,
                                                          down_expert_bytes,
                                                          force_reuse,
                                                          &gate_buf,
                                                          &up_buf,
                                                          &down_buf,
                                                          &gate_inner,
                                                          &up_inner,
                                                          &down_inner)) {
        return NULL;
    }
    if (!gate_buf || !up_buf || !down_buf) return NULL;

    uint8_t *gate_dst = (uint8_t *)[gate_buf contents] + gate_inner;
    uint8_t *up_dst = (uint8_t *)[up_buf contents] + up_inner;
    uint8_t *down_dst = (uint8_t *)[down_buf contents] + down_inner;
    if (!gate_dst || !up_dst || !down_dst) return NULL;

    ds4_gpu_stream_expert_pread_task tasks[3] = {
        {
            .offset = gate_abs_offset,
            .len = gate_expert_bytes,
            .dst = gate_dst,
        },
        {
            .offset = up_abs_offset,
            .len = gate_expert_bytes,
            .dst = up_dst,
        },
        {
            .offset = down_abs_offset,
            .len = down_expert_bytes,
            .dst = down_dst,
        },
    };
    uint64_t read_bytes = 0;
    double read_ms = 0.0;
    if (!ds4_gpu_stream_expert_pread_tasks(tasks, 3, &read_bytes, &read_ms)) {
        return NULL;
    }
    ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms);

    [gate_buf didModifyRange:NSMakeRange(gate_inner, (NSUInteger)gate_expert_bytes)];
    [up_buf didModifyRange:NSMakeRange(up_inner, (NSUInteger)gate_expert_bytes)];
    [down_buf didModifyRange:NSMakeRange(down_inner, (NSUInteger)down_expert_bytes)];
    if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) {
        fprintf(stderr,
                "ds4: Metal streaming expert parallel pread layer=%u experts=1 tensors=3 "
                "threads=%u bytes=%.2f GiB wall=%.3f ms\n",
                layer,
                ds4_gpu_stream_expert_pread_thread_count(3),
                ds4_gpu_gib(read_bytes),
                read_ms);
    }
    return ds4_gpu_stream_expert_cache_install_loaded(model_map,
                                                      model_size,
                                                      layer,
                                                      expert,
                                                      gate_abs_offset,
                                                      up_abs_offset,
                                                      down_abs_offset,
                                                      gate_expert_bytes,
                                                      down_expert_bytes,
                                                      gate_buf,
                                                      up_buf,
                                                      down_buf,
                                                      gate_inner,
                                                      up_inner,
                                                      down_inner);
}

static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get(
        const void *model_map,
        uint64_t    model_size,
        uint32_t    layer,
        uint32_t    expert,
        uint32_t    n_total_expert,
        uint32_t    n_selected,
        uint64_t    gate_abs_offset,
        uint64_t    up_abs_offset,
        uint64_t    down_abs_offset,
        uint64_t    gate_expert_bytes,
        uint64_t    down_expert_bytes) {
    return ds4_gpu_stream_expert_cache_get_protected(model_map,
                                                     model_size,
                                                     layer,
                                                     expert,
                                                     n_total_expert,
                                                     n_selected,
                                                     gate_abs_offset,
                                                     up_abs_offset,
                                                     down_abs_offset,
                                                     gate_expert_bytes,
                                                     down_expert_bytes,
                                                     NULL,
                                                     0);
}

static int ds4_gpu_stream_expert_pending_load_profile_enabled(void) {
    return getenv("DS4_METAL_STREAMING_EXPERT_EARLY_LOAD_PROFILE") != NULL;
}

static void ds4_gpu_stream_expert_pending_load_release_buffers(
        ds4_gpu_stream_expert_pending_load *p) {
    if (!p) return;
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        p->gate_bufs[i] = nil;
        p->up_bufs[i] = nil;
        p->down_bufs[i] = nil;
        p->gate_inners[i] = 0;
        p->up_inners[i] = 0;
        p->down_inners[i] = 0;
    }
}

static int ds4_gpu_stream_expert_pending_load_install(
        ds4_gpu_stream_expert_pending_load *p,
        ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED],
        double elapsed_ms) {
    if (!p || p->n_loads == 0) return 1;

    uint64_t read_bytes = 0;
    int ok = 1;
    for (uint32_t i = 0; i < p->n_tasks; i++) {
        if (!p->tasks[i].ok) ok = 0;
        if (read_bytes > UINT64_MAX - p->tasks[i].read_bytes) {
            read_bytes = UINT64_MAX;
        } else {
            read_bytes += p->tasks[i].read_bytes;
        }
    }
    if (!ok) return 0;

    ds4_gpu_stream_expert_cache_note_pread(p->layer, read_bytes, elapsed_ms);
    const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled();
    double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
    double load_modify_ms = 0.0;
    double load_install_ms = 0.0;
    for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) {
        [p->gate_bufs[load_i] didModifyRange:NSMakeRange(p->gate_inners[load_i], (NSUInteger)p->gate_expert_bytes)];
        [p->up_bufs[load_i] didModifyRange:NSMakeRange(p->up_inners[load_i], (NSUInteger)p->gate_expert_bytes)];
        [p->down_bufs[load_i] didModifyRange:NSMakeRange(p->down_inners[load_i], (NSUInteger)p->down_expert_bytes)];
    }
    if (load_timing) {
        const double now_ms = ds4_gpu_now_ms();
        load_modify_ms = now_ms - load_t0;
        load_t0 = now_ms;
    }

    ds4_gpu_stream_expert_cache_entry
        *loaded_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        loaded_entries[i] = NULL;
    }
    for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) {
        const uint32_t slot = p->load_slots[load_i];
        const uint32_t expert = (uint32_t)p->selected_ids[slot];
        ds4_gpu_stream_expert_cache_entry *entry =
            ds4_gpu_stream_expert_cache_install_loaded(p->model_map,
                                                       p->model_size,
                                                       p->layer,
                                                       expert,
                                                       p->gate_abs_offsets[slot],
                                                       p->up_abs_offsets[slot],
                                                       p->down_abs_offsets[slot],
                                                       p->gate_expert_bytes,
                                                       p->down_expert_bytes,
                                                       p->gate_bufs[load_i],
                                                       p->up_bufs[load_i],
                                                       p->down_bufs[load_i],
                                                       p->gate_inners[load_i],
                                                       p->up_inners[load_i],
                                                       p->down_inners[load_i]);
        if (!entry) return 0;
        loaded_entries[slot] = entry;
        if (entries) entries[slot] = entry;
    }
    for (uint32_t i = 0; i < p->n_selected; i++) {
        if ((p->missing_mask & (1u << i)) == 0) continue;
        if (entries && entries[i]) continue;
        const uint32_t source = p->source_slots[i];
        if (source >= p->n_selected) return 0;
        ds4_gpu_stream_expert_cache_entry *entry = entries && entries[source] ?
            entries[source] : loaded_entries[source];
        if (!entry) return 0;
        entry->use_count++;
        if (entries) entries[i] = entry;
    }
    if (load_timing) {
        load_install_ms = ds4_gpu_now_ms() - load_t0;
        ds4_gpu_stream_expert_timing_note_load_detail(p->prepare_ms,
                                                      elapsed_ms,
                                                      load_modify_ms,
                                                      load_install_ms);
    }
    if (ds4_gpu_stream_expert_pending_load_profile_enabled()) {
        fprintf(stderr,
                "ds4: Metal streaming expert early-load finish layer=%u experts=%u tensors=%u bytes=%.2f GiB wall=%.3f ms\n",
                p->layer,
                p->n_loads,
                p->n_tasks,
                ds4_gpu_gib(read_bytes),
                elapsed_ms);
    }
    return 1;
}

static int ds4_gpu_stream_expert_pending_load_finish(
        ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]) {
    ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load;
    if (!p->active) return 1;

    const double start_ms = p->start_ms;
    if (!ds4_gpu_stream_expert_pread_pool_wait()) {
        ds4_gpu_stream_expert_pending_load_release_buffers(p);
        p->active = 0;
        return 0;
    }
    const double elapsed_ms = ds4_gpu_now_ms() - start_ms;
    p->active = 0;
    const int ok = ds4_gpu_stream_expert_pending_load_install(p,
                                                              entries,
                                                              elapsed_ms);
    ds4_gpu_stream_expert_pending_load_release_buffers(p);
    p->n_tasks = 0;
    p->n_loads = 0;
    p->prepare_ms = 0.0;
    return ok;
}

static void ds4_gpu_stream_expert_pending_load_clear(void) {
    if (!g_stream_expert_pending_load.active) {
        ds4_gpu_stream_expert_pending_load_release_buffers(
                &g_stream_expert_pending_load);
        return;
    }
    (void)ds4_gpu_stream_expert_pending_load_finish(NULL);
}

static int ds4_gpu_stream_expert_pending_load_matches(
        const void   *model_map,
        uint64_t      model_size,
        uint32_t      layer,
        const int32_t *selected_ids,
        uint32_t      n_total_expert,
        uint32_t      n_selected,
        uint64_t      gate_expert_bytes,
        uint64_t      down_expert_bytes) {
    ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load;
    if (!p->active ||
        p->model_map != model_map ||
        p->model_size != model_size ||
        p->layer != layer ||
        p->n_total_expert != n_total_expert ||
        p->n_selected != n_selected ||
        p->gate_expert_bytes != gate_expert_bytes ||
        p->down_expert_bytes != down_expert_bytes) {
        return 0;
    }
    for (uint32_t i = 0; i < n_selected; i++) {
        if (p->selected_ids[i] != selected_ids[i]) return 0;
    }
    return 1;
}

int ds4_gpu_stream_expert_cache_begin_selected_load(
        const ds4_gpu_stream_expert_table *table,
        const int32_t                     *selected_ids,
        uint32_t                           n_selected) {
    if (!g_ssd_streaming_mode ||
        getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EARLY_LOAD") != NULL) {
        return 1;
    }
    if (!table) return 0;
    const void *model_map = table->model_map;
    const uint64_t model_size = table->model_size;
    const uint32_t layer = table->layer;
    const uint32_t n_total_expert = table->n_total_expert;
    const uint64_t gate_offset = table->gate_offset;
    const uint64_t up_offset = table->up_offset;
    const uint64_t down_offset = table->down_offset;
    const uint64_t gate_expert_bytes = table->gate_expert_bytes;
    const uint64_t down_expert_bytes = table->down_expert_bytes;
    if (!model_map || !selected_ids ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_selected == 0 ||
        n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
        n_total_expert == 0 ||
        n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes) ||
        ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                  n_total_expert,
                                                  n_selected) == 0) {
        return 1;
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;

    if (ds4_gpu_stream_expert_pending_load_matches(model_map,
                                                   model_size,
                                                   layer,
                                                   selected_ids,
                                                   n_total_expert,
                                                   n_selected,
                                                   gate_expert_bytes,
                                                   down_expert_bytes)) {
        return 1;
    }

    ds4_gpu_stream_expert_pending_load_clear();
    ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load;
    p->active = 0;
    p->model_map = model_map;
    p->model_size = model_size;
    p->layer = layer;
    p->n_total_expert = n_total_expert;
    p->n_selected = n_selected;
    p->missing_mask = 0;
    p->n_loads = 0;
    p->n_tasks = 0;
    p->gate_expert_bytes = gate_expert_bytes;
    p->down_expert_bytes = down_expert_bytes;
    p->prepare_ms = 0.0;
    const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled();
    const double load_prepare_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        p->selected_ids[i] = -1;
        p->load_slots[i] = 0;
        p->source_slots[i] = UINT32_MAX;
        p->gate_abs_offsets[i] = 0;
        p->up_abs_offsets[i] = 0;
        p->down_abs_offsets[i] = 0;
        p->gate_bufs[i] = nil;
        p->up_bufs[i] = nil;
        p->down_bufs[i] = nil;
        p->gate_inners[i] = 0;
        p->up_inners[i] = 0;
        p->down_inners[i] = 0;
    }
    memset(p->tasks, 0, sizeof(p->tasks));

    for (uint32_t i = 0; i < n_selected; i++) {
        if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
            fprintf(stderr,
                    "ds4: Metal streaming early-load expert id %d is outside 0..%u\n",
                    selected_ids[i],
                    n_total_expert);
            return 0;
        }
        p->selected_ids[i] = selected_ids[i];
        const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i];
        if (expert_id > UINT64_MAX / gate_expert_bytes ||
            expert_id > UINT64_MAX / down_expert_bytes) {
            fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n");
            return 0;
        }
        const uint64_t gate_rel = expert_id * gate_expert_bytes;
        const uint64_t down_rel = expert_id * down_expert_bytes;
        if (gate_rel > UINT64_MAX - gate_offset ||
            gate_rel > UINT64_MAX - up_offset ||
            down_rel > UINT64_MAX - down_offset) {
            fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n");
            return 0;
        }
        p->gate_abs_offsets[i] = gate_offset + gate_rel;
        p->up_abs_offsets[i] = up_offset + gate_rel;
        p->down_abs_offsets[i] = down_offset + down_rel;

        ds4_gpu_stream_expert_cache_entry *e =
            &g_stream_expert_cache[layer][(uint32_t)selected_ids[i]];
        if (ds4_gpu_stream_expert_cache_entry_matches(e,
                                                      model_map,
                                                      model_size,
                                                      p->gate_abs_offsets[i],
                                                      p->up_abs_offsets[i],
                                                      p->down_abs_offsets[i],
                                                      gate_expert_bytes,
                                                      down_expert_bytes)) {
            continue;
        }

        uint32_t source = UINT32_MAX;
        for (uint32_t prev = 0; prev < i; prev++) {
            if (selected_ids[prev] == selected_ids[i] &&
                p->gate_abs_offsets[prev] == p->gate_abs_offsets[i] &&
                p->up_abs_offsets[prev] == p->up_abs_offsets[i] &&
                p->down_abs_offsets[prev] == p->down_abs_offsets[i] &&
                (p->missing_mask & (1u << prev)) != 0) {
                source = p->source_slots[prev] != UINT32_MAX ?
                         p->source_slots[prev] : prev;
                break;
            }
        }
        if (source != UINT32_MAX) {
            p->source_slots[i] = source;
            p->missing_mask |= 1u << i;
            continue;
        }
        p->source_slots[i] = i;
        p->missing_mask |= 1u << i;
        p->load_slots[p->n_loads++] = i;
    }
    if (p->n_loads == 0) return 1;

    const uint32_t cache_budget =
        ds4_gpu_stream_expert_cache_configured_budget();
    uint32_t reserved_entries = g_stream_expert_cache_entry_count;
    ds4_gpu_stream_expert_reusable_buffers
        batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        batch_reuse[i] =
            (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 };
    }
    uint32_t batch_reuse_count = 0;
    if (cache_budget != 0 &&
        reserved_entries >= cache_budget &&
        p->n_loads > 1 &&
        ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes,
                                                  down_expert_bytes)) {
        const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
        batch_reuse_count =
            ds4_gpu_stream_expert_cache_take_reusable_batch(
                    p->n_loads,
                    layer,
                    selected_ids,
                    n_selected,
                    gate_expert_bytes,
                    down_expert_bytes,
                    batch_reuse);
        if (load_timing) {
            ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(
                    ds4_gpu_now_ms() - reuse_t0);
        }
    }
    for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) {
        const uint32_t slot = p->load_slots[load_i];
        const uint32_t expert = (uint32_t)p->selected_ids[slot];
        const int force_reuse =
            cache_budget != 0 && reserved_entries >= cache_budget;

        ds4_gpu_stream_expert_readahead_range(p->gate_abs_offsets[slot],
                                              gate_expert_bytes);
        ds4_gpu_stream_expert_readahead_range(p->up_abs_offsets[slot],
                                              gate_expert_bytes);
        ds4_gpu_stream_expert_readahead_range(p->down_abs_offsets[slot],
                                              down_expert_bytes);

        if (load_i < batch_reuse_count &&
            batch_reuse[load_i].gate_buffer &&
            batch_reuse[load_i].up_buffer &&
            batch_reuse[load_i].down_buffer) {
            p->gate_bufs[load_i] = batch_reuse[load_i].gate_buffer;
            p->up_bufs[load_i] = batch_reuse[load_i].up_buffer;
            p->down_bufs[load_i] = batch_reuse[load_i].down_buffer;
            p->gate_inners[load_i] = batch_reuse[load_i].gate_inner;
            p->up_inners[load_i] = batch_reuse[load_i].up_inner;
            p->down_inners[load_i] = batch_reuse[load_i].down_inner;
        } else {
            const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
            const int prepared =
                ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
                                                                 expert,
                                                                 layer,
                                                                 selected_ids,
                                                                 n_selected,
                                                                 gate_expert_bytes,
                                                                 down_expert_bytes,
                                                                 force_reuse,
                                                                 &p->gate_bufs[load_i],
                                                                 &p->up_bufs[load_i],
                                                                 &p->down_bufs[load_i],
                                                                 &p->gate_inners[load_i],
                                                                 &p->up_inners[load_i],
                                                                 &p->down_inners[load_i]);
            if (load_timing) {
                ds4_gpu_stream_expert_timing_note_prepare_buffer(
                        ds4_gpu_now_ms() - buffer_t0);
            }
            if (!prepared) {
                ds4_gpu_stream_expert_pending_load_release_buffers(p);
                return 0;
            }
        }
        if (!force_reuse && reserved_entries < UINT32_MAX) {
            reserved_entries++;
        }
        uint8_t *gate_dst = (uint8_t *)[p->gate_bufs[load_i] contents] +
                             p->gate_inners[load_i];
        uint8_t *up_dst = (uint8_t *)[p->up_bufs[load_i] contents] +
                           p->up_inners[load_i];
        uint8_t *down_dst = (uint8_t *)[p->down_bufs[load_i] contents] +
                             p->down_inners[load_i];
        if (!gate_dst || !up_dst || !down_dst) {
            ds4_gpu_stream_expert_pending_load_release_buffers(p);
            return 0;
        }
        const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
        p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
            .offset = p->gate_abs_offsets[slot],
            .len = gate_expert_bytes,
            .dst = gate_dst,
        };
        p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
            .offset = p->up_abs_offsets[slot],
            .len = gate_expert_bytes,
            .dst = up_dst,
        };
        p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
            .offset = p->down_abs_offsets[slot],
            .len = down_expert_bytes,
            .dst = down_dst,
        };
        if (load_timing) {
            ds4_gpu_stream_expert_timing_note_prepare_task(
                    1,
                    ds4_gpu_now_ms() - task_t0);
        }
    }

    const uint32_t n_workers =
        ds4_gpu_stream_expert_pread_thread_count(p->n_tasks);
    p->start_ms = ds4_gpu_now_ms();
    if (load_timing) {
        p->prepare_ms = p->start_ms - load_prepare_t0;
    }
    if (ds4_gpu_stream_expert_pread_pool_begin(p->tasks,
                                               p->n_tasks,
                                               n_workers)) {
        p->active = 1;
        if (ds4_gpu_stream_expert_pending_load_profile_enabled()) {
            fprintf(stderr,
                    "ds4: Metal streaming expert early-load begin layer=%u experts=%u tensors=%u threads=%u\n",
                    layer,
                    p->n_loads,
                    p->n_tasks,
                    n_workers);
        }
        return 1;
    }

    uint64_t read_bytes = 0;
    double read_ms = 0.0;
    if (!ds4_gpu_stream_expert_pread_tasks(p->tasks,
                                           p->n_tasks,
                                           &read_bytes,
                                           &read_ms)) {
        ds4_gpu_stream_expert_pending_load_release_buffers(p);
        return 0;
    }
    (void)read_bytes;
    if (!ds4_gpu_stream_expert_pending_load_install(p, NULL, read_ms)) {
        ds4_gpu_stream_expert_pending_load_release_buffers(p);
        return 0;
    }
    ds4_gpu_stream_expert_pending_load_release_buffers(p);
    p->n_tasks = 0;
    p->n_loads = 0;
    p->prepare_ms = 0.0;
    return 1;
}

static int ds4_gpu_stream_expert_cache_load_selected_missing(
        const void    *model_map,
        uint64_t       model_size,
        uint32_t       layer,
        const int32_t *selected_ids,
        uint32_t       n_total_expert,
        uint32_t       n_selected,
        const uint64_t *gate_abs_offsets,
        const uint64_t *up_abs_offsets,
        const uint64_t *down_abs_offsets,
        uint64_t       gate_expert_bytes,
        uint64_t       down_expert_bytes,
        uint32_t       missing_mask,
        ds4_gpu_stream_expert_cache_entry **entries) {
    if (!g_ssd_streaming_mode ||
        !model_map ||
        !selected_ids ||
        !gate_abs_offsets ||
        !up_abs_offsets ||
        !down_abs_offsets ||
        !entries ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_selected == 0 ||
        n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
        n_total_expert == 0 ||
        n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes) ||
        ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                  n_total_expert,
                                                  n_selected) == 0) {
        return 0;
    }
    missing_mask &= (1u << n_selected) - 1u;
    if (missing_mask == 0) return 1;
    if (g_stream_expert_pending_load.active &&
        n_selected <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
        if (ds4_gpu_stream_expert_pending_load_matches(model_map,
                                                       model_size,
                                                       layer,
                                                       selected_ids,
                                                       n_total_expert,
                                                       n_selected,
                                                       gate_expert_bytes,
                                                       down_expert_bytes)) {
            if (!ds4_gpu_stream_expert_pending_load_finish(entries)) return 0;
            for (uint32_t i = 0; i < n_selected; i++) {
                if ((missing_mask & (1u << i)) != 0 && entries[i]) {
                    missing_mask &= ~(1u << i);
                }
            }
            if (missing_mask == 0) return 1;
        } else {
            ds4_gpu_stream_expert_pending_load_clear();
        }
    } else if (g_stream_expert_pending_load.active) {
        ds4_gpu_stream_expert_pending_load_clear();
    }

    const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled();
    double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
    double load_prepare_ms = 0.0;
    double load_modify_ms = 0.0;
    double load_install_ms = 0.0;

    uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        load_slots[i] = 0;
        source_slots[i] = UINT32_MAX;
    }
    uint32_t n_loads = 0;
    for (uint32_t i = 0; i < n_selected; i++) {
        if ((missing_mask & (1u << i)) == 0) continue;
        if (entries[i]) {
            source_slots[i] = i;
            continue;
        }
        if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
            fprintf(stderr,
                    "ds4: Metal streaming selected missing expert id %d is outside 0..%u\n",
                    selected_ids[i],
                    n_total_expert);
            return 0;
        }
        uint32_t source = UINT32_MAX;
        for (uint32_t prev = 0; prev < i; prev++) {
            if (selected_ids[prev] == selected_ids[i] &&
                gate_abs_offsets[prev] == gate_abs_offsets[i] &&
                up_abs_offsets[prev] == up_abs_offsets[i] &&
                down_abs_offsets[prev] == down_abs_offsets[i] &&
                (entries[prev] || (missing_mask & (1u << prev)) != 0)) {
                source = source_slots[prev] != UINT32_MAX ?
                         source_slots[prev] : prev;
                break;
            }
        }
        if (source != UINT32_MAX) {
            source_slots[i] = source;
            continue;
        }
        source_slots[i] = i;
        load_slots[n_loads++] = i;
    }
    if (n_loads == 0) {
        for (uint32_t i = 0; i < n_selected; i++) {
            if ((missing_mask & (1u << i)) == 0 || entries[i]) continue;
            const uint32_t source = source_slots[i];
            if (source >= n_selected || !entries[source]) return 0;
            entries[i] = entries[source];
            entries[i]->use_count++;
        }
        return 1;
    }

    __strong id<MTLBuffer> gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    __strong id<MTLBuffer> up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    __strong id<MTLBuffer> down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        gate_bufs[i] = nil;
        up_bufs[i] = nil;
        down_bufs[i] = nil;
        gate_inners[i] = 0;
        up_inners[i] = 0;
        down_inners[i] = 0;
    }
    ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u];
    memset(tasks, 0, sizeof(tasks));
    uint32_t n_tasks = 0;
    const uint32_t cache_budget =
        ds4_gpu_stream_expert_cache_configured_budget();
    uint32_t reserved_entries = g_stream_expert_cache_entry_count;
    ds4_gpu_stream_expert_reusable_buffers
        batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        batch_reuse[i] =
            (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 };
    }
    uint32_t batch_reuse_count = 0;
    if (cache_budget != 0 &&
        reserved_entries >= cache_budget &&
        n_loads > 1 &&
        ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes,
                                                  down_expert_bytes)) {
        const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
        batch_reuse_count =
            ds4_gpu_stream_expert_cache_take_reusable_batch(
                    n_loads,
                    layer,
                    selected_ids,
                    n_selected,
                    gate_expert_bytes,
                    down_expert_bytes,
                    batch_reuse);
        if (load_timing) {
            ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(
                    ds4_gpu_now_ms() - reuse_t0);
        }
    }

    for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
        const uint32_t slot = load_slots[load_i];
        const uint32_t expert = (uint32_t)selected_ids[slot];
        const int force_reuse =
            cache_budget != 0 && reserved_entries >= cache_budget;

        ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot],
                                              gate_expert_bytes);
        ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot],
                                              gate_expert_bytes);
        ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot],
                                              down_expert_bytes);

        if (load_i < batch_reuse_count &&
            batch_reuse[load_i].gate_buffer &&
            batch_reuse[load_i].up_buffer &&
            batch_reuse[load_i].down_buffer) {
            gate_bufs[load_i] = batch_reuse[load_i].gate_buffer;
            up_bufs[load_i] = batch_reuse[load_i].up_buffer;
            down_bufs[load_i] = batch_reuse[load_i].down_buffer;
            gate_inners[load_i] = batch_reuse[load_i].gate_inner;
            up_inners[load_i] = batch_reuse[load_i].up_inner;
            down_inners[load_i] = batch_reuse[load_i].down_inner;
        } else {
            const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
            const int prepared =
                ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
                                                                 expert,
                                                                 layer,
                                                                 selected_ids,
                                                                 n_selected,
                                                                 gate_expert_bytes,
                                                                 down_expert_bytes,
                                                                 force_reuse,
                                                                 &gate_bufs[load_i],
                                                                 &up_bufs[load_i],
                                                                 &down_bufs[load_i],
                                                                 &gate_inners[load_i],
                                                                 &up_inners[load_i],
                                                                 &down_inners[load_i]);
            if (load_timing) {
                ds4_gpu_stream_expert_timing_note_prepare_buffer(
                        ds4_gpu_now_ms() - buffer_t0);
            }
            if (!prepared) {
                return 0;
            }
        }
        if (!force_reuse && reserved_entries < UINT32_MAX) {
            reserved_entries++;
        }
        if (!gate_bufs[load_i] || !up_bufs[load_i] || !down_bufs[load_i]) {
            return 0;
        }

        uint8_t *gate_dst = (uint8_t *)[gate_bufs[load_i] contents] +
                             gate_inners[load_i];
        uint8_t *up_dst = (uint8_t *)[up_bufs[load_i] contents] +
                           up_inners[load_i];
        uint8_t *down_dst = (uint8_t *)[down_bufs[load_i] contents] +
                             down_inners[load_i];
        if (!gate_dst || !up_dst || !down_dst) return 0;

        const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
        tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
            .offset = gate_abs_offsets[slot],
            .len = gate_expert_bytes,
            .dst = gate_dst,
        };
        tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
            .offset = up_abs_offsets[slot],
            .len = gate_expert_bytes,
            .dst = up_dst,
        };
        tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
            .offset = down_abs_offsets[slot],
            .len = down_expert_bytes,
            .dst = down_dst,
        };
        if (load_timing) {
            ds4_gpu_stream_expert_timing_note_prepare_task(
                    1,
                    ds4_gpu_now_ms() - task_t0);
        }
    }

    if (load_timing) {
        const double now_ms = ds4_gpu_now_ms();
        load_prepare_ms = now_ms - load_t0;
        load_t0 = now_ms;
    }

    uint64_t read_bytes = 0;
    double read_ms = 0.0;
    const int ok = ds4_gpu_stream_expert_pread_tasks(tasks,
                                                     n_tasks,
                                                     &read_bytes,
                                                     &read_ms);
    if (!ok) return 0;
    if (load_timing) {
        load_t0 = ds4_gpu_now_ms();
    }
    ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms);

    for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
        [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)];
        [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)];
        [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)];
    }
    if (load_timing) {
        const double now_ms = ds4_gpu_now_ms();
        load_modify_ms = now_ms - load_t0;
        load_t0 = now_ms;
    }
    if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) {
        fprintf(stderr,
                "ds4: Metal streaming expert parallel pread layer=%u experts=%u tensors=%u "
                "threads=%u bytes=%.2f GiB wall=%.3f ms\n",
                layer,
                n_loads,
                n_tasks,
                ds4_gpu_stream_expert_pread_thread_count(n_tasks),
                ds4_gpu_gib(read_bytes),
                read_ms);
    }

    for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
        const uint32_t slot = load_slots[load_i];
        const uint32_t expert = (uint32_t)selected_ids[slot];
        ds4_gpu_stream_expert_cache_entry *entry =
            ds4_gpu_stream_expert_cache_install_loaded(model_map,
                                                       model_size,
                                                       layer,
                                                       expert,
                                                       gate_abs_offsets[slot],
                                                       up_abs_offsets[slot],
                                                       down_abs_offsets[slot],
                                                       gate_expert_bytes,
                                                       down_expert_bytes,
                                                       gate_bufs[load_i],
                                                       up_bufs[load_i],
                                                       down_bufs[load_i],
                                                       gate_inners[load_i],
                                                       up_inners[load_i],
                                                       down_inners[load_i]);
        if (!entry) return 0;
        entries[slot] = entry;
    }

    for (uint32_t i = 0; i < n_selected; i++) {
        if ((missing_mask & (1u << i)) == 0 || entries[i]) continue;
        const uint32_t source = source_slots[i];
        if (source >= n_selected || !entries[source]) return 0;
        entries[i] = entries[source];
        entries[i]->use_count++;
    }
    for (uint32_t i = 0; i < n_selected; i++) {
        if ((missing_mask & (1u << i)) != 0 && !entries[i]) return 0;
    }
    if (load_timing) {
        load_install_ms = ds4_gpu_now_ms() - load_t0;
        ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms,
                                                      read_ms,
                                                      load_modify_ms,
                                                      load_install_ms);
    }
    return 1;
}

static void ds4_gpu_glm_stream_selected_prefetch_set(
        const ds4_gpu_stream_expert_table *table,
        const int32_t                     *selected_ids,
        uint32_t                           n_selected) {
    g_glm_stream_selected_prefetch.active = 0;
    if (!table || !selected_ids ||
        n_selected == 0 ||
        n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
        return;
    }
    g_glm_stream_selected_prefetch.model_map = table->model_map;
    g_glm_stream_selected_prefetch.model_size = table->model_size;
    g_glm_stream_selected_prefetch.layer = table->layer;
    g_glm_stream_selected_prefetch.n_total_expert = table->n_total_expert;
    g_glm_stream_selected_prefetch.n_selected = n_selected;
    g_glm_stream_selected_prefetch.gate_offset = table->gate_offset;
    g_glm_stream_selected_prefetch.up_offset = table->up_offset;
    g_glm_stream_selected_prefetch.down_offset = table->down_offset;
    g_glm_stream_selected_prefetch.gate_expert_bytes = table->gate_expert_bytes;
    g_glm_stream_selected_prefetch.down_expert_bytes = table->down_expert_bytes;
    for (uint32_t i = 0; i < n_selected; i++) {
        g_glm_stream_selected_prefetch.selected_ids[i] = selected_ids[i];
    }
    g_glm_stream_selected_prefetch.active = 1;
}

static int ds4_gpu_glm_stream_selected_prefetch_take(
        const void *model_map,
        uint64_t    model_size,
        uint32_t    layer,
        uint32_t    n_total_expert,
        uint32_t    n_selected,
        uint64_t    gate_offset,
        uint64_t    up_offset,
        uint64_t    down_offset,
        uint64_t    gate_expert_bytes,
        uint64_t    down_expert_bytes,
        int32_t    *selected_ids_out) {
    if (!selected_ids_out || !g_glm_stream_selected_prefetch.active) return 0;
    ds4_gpu_glm_stream_selected_prefetch *p = &g_glm_stream_selected_prefetch;
    if (p->model_map != model_map ||
        p->model_size != model_size ||
        p->layer != layer ||
        p->n_total_expert != n_total_expert ||
        p->n_selected != n_selected ||
        p->gate_offset != gate_offset ||
        p->up_offset != up_offset ||
        p->down_offset != down_offset ||
        p->gate_expert_bytes != gate_expert_bytes ||
        p->down_expert_bytes != down_expert_bytes) {
        return 0;
    }
    for (uint32_t i = 0; i < n_selected; i++) {
        selected_ids_out[i] = p->selected_ids[i];
    }
    p->active = 0;
    return 1;
}

int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
        const ds4_gpu_stream_expert_table *table,
        const ds4_gpu_tensor              *selected,
        uint32_t                           n_selected) {
    g_glm_stream_selected_prefetch.active = 0;
    if (!g_ssd_streaming_mode ||
        getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL) {
        return 1;
    }
    if (!table || !selected ||
        n_selected == 0 ||
        n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
        return 1;
    }

    const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
    double t0 = timing ? ds4_gpu_now_ms() : 0.0;
    double selected_sync_ms = 0.0;
    double selected_copy_ms = 0.0;
    double selected_bind_ms = 0.0;
    const int had_batch = g_batch_cb != nil;
    if (had_batch) {
        if (ds4_gpu_end_commands() == 0) return 0;
        if (timing) {
            selected_sync_ms = ds4_gpu_now_ms() - t0;
            t0 = ds4_gpu_now_ms();
        }
    }

    int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
        selected_ids[i] = -1;
    }
    int ok = ds4_gpu_tensor_read(selected,
                                 0,
                                 selected_ids,
                                 (uint64_t)n_selected * sizeof(selected_ids[0])) != 0;
    if (timing) {
        selected_copy_ms = ds4_gpu_now_ms() - t0;
        t0 = ds4_gpu_now_ms();
    }
    if (ok) {
        ok = ds4_gpu_stream_expert_cache_begin_selected_load(table,
                                                             selected_ids,
                                                             n_selected) != 0;
    }
    if (timing) {
        selected_bind_ms = ds4_gpu_now_ms() - t0;
        ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms,
                                                   selected_copy_ms,
                                                   selected_bind_ms);
    }
    if (ok) {
        ds4_gpu_glm_stream_selected_prefetch_set(table,
                                                 selected_ids,
                                                 n_selected);
    }
    if (had_batch && ds4_gpu_begin_commands() == 0) ok = 0;
    return ok;
}

static void ds4_gpu_stream_expert_cache_clear_layer(uint32_t layer) {
    if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return;
    for (uint32_t expert = 0;
         expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
         expert++) {
        ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
    }
    g_stream_expert_cache_layer_count[layer] = 0;
}

static int ds4_gpu_stream_expert_cache_prepare_selected_batch(
        const void    *model_map,
        uint64_t       model_size,
        uint32_t       layer,
        const ds4_gpu_tensor *selected,
        uint32_t       n_tokens,
        uint32_t       n_total_expert,
        uint32_t       n_selected,
        uint64_t       gate_offset,
        uint64_t       up_offset,
        uint64_t       down_offset,
        uint64_t       gate_expert_bytes,
        uint64_t       down_expert_bytes,
        id<MTLBuffer> *gate_addrs,
        id<MTLBuffer> *up_addrs,
        id<MTLBuffer> *down_addrs,
        ds4_gpu_stream_expert_cache_entry **resources,
        uint32_t      *n_resources,
        uint32_t      *unique_out,
        id<MTLBuffer> *overflow_gate,
        id<MTLBuffer> *overflow_up,
        id<MTLBuffer> *overflow_down) {
    if (overflow_gate) *overflow_gate = nil;
    if (overflow_up) *overflow_up = nil;
    if (overflow_down) *overflow_down = nil;
    if (!g_ssd_streaming_mode ||
        !model_map ||
        !selected ||
        !gate_addrs ||
        !up_addrs ||
        !down_addrs ||
        !resources ||
        !n_resources ||
        !unique_out ||
        !overflow_gate ||
        !overflow_up ||
        !overflow_down ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_tokens == 0 ||
        n_selected == 0 ||
        n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
        n_total_expert == 0 ||
        n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        gate_expert_bytes == 0 ||
        down_expert_bytes == 0) {
        return 0;
    }
    if (n_tokens > UINT32_MAX / n_selected) return 0;

    const uint64_t n_ids = (uint64_t)n_tokens * n_selected;
    if (n_ids > SIZE_MAX / sizeof(int32_t)) return 0;
    int32_t *ids = malloc((size_t)n_ids * sizeof(ids[0]));
    if (!ids) return 0;

    const bool profile =
        getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_PROFILE") != NULL;
    const double t0 = profile ? ds4_gpu_now_ms() : 0.0;
    int ok = ds4_gpu_tensor_read(selected,
                                 0,
                                 ids,
                                 n_ids * sizeof(ids[0]));
    const double t_read = profile ? ds4_gpu_now_ms() : 0.0;

    bool seen[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { false };
    uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { 0 };
    int32_t unique_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    uint32_t unique_count = 0;
    *n_resources = 0;
    *unique_out = 0;
    if (ok) {
        if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) {
            ok = 0;
        }
    }
    if (ok) {
        for (uint64_t i = 0; i < n_ids; i++) {
            const int32_t selected_id = ids[i];
            if (selected_id < 0 || (uint32_t)selected_id >= n_total_expert) {
                fprintf(stderr,
                        "ds4: Metal streaming batch selected expert id %d is outside 0..%u at layer %u\n",
                        selected_id,
                        n_total_expert,
                        layer);
                ok = 0;
                break;
            }
            frequency[(uint32_t)selected_id]++;
            if (!seen[(uint32_t)selected_id]) {
                seen[(uint32_t)selected_id] = true;
                unique_ids[unique_count++] = selected_id;
            }
        }
    }
    if (ok) {
        ds4_gpu_stream_expert_cache_note_frequency_hotness(layer,
                                                           frequency,
                                                           n_total_expert);
    }
    /*
     * When the layer's unique selected set does not fit the cache budget, the
     * extra experts are addressed straight into whole-tensor mapped model views
     * instead of falling back to a different MoE kernel path. The address-table
     * kernels read identical expert bytes either way, so the cache/view split
     * does not change the computed logits.
     */
    uint32_t view_served = 0;
    uint64_t overflow_gate_inner = 0;
    uint64_t overflow_up_inner = 0;
    uint64_t overflow_down_inner = 0;

    uint64_t unique_gate_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    uint64_t unique_up_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    uint64_t unique_down_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    ds4_gpu_stream_expert_cache_entry
        *unique_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    uint32_t load_unique[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    __strong id<MTLBuffer>
        gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    __strong id<MTLBuffer>
        up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    __strong id<MTLBuffer>
        down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
    for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; i++) {
        unique_gate_offsets[i] = 0;
        unique_up_offsets[i] = 0;
        unique_down_offsets[i] = 0;
        unique_entries[i] = NULL;
        load_unique[i] = 0;
        gate_bufs[i] = nil;
        up_bufs[i] = nil;
        down_bufs[i] = nil;
        gate_inners[i] = 0;
        up_inners[i] = 0;
        down_inners[i] = 0;
    }

    ds4_gpu_stream_expert_pread_task *tasks = NULL;
    uint32_t n_loads = 0;
    uint32_t n_tasks = 0;
    double load_prepare_ms = 0.0;
    double load_modify_ms = 0.0;
    double load_install_ms = 0.0;
    double load_timing_t0 = ds4_gpu_stream_expert_timing_summary_enabled() ?
                            ds4_gpu_now_ms() : 0.0;
    const int load_timing = load_timing_t0 != 0.0;
    if (ok && unique_count != 0) {
        tasks = calloc((size_t)unique_count * 3u, sizeof(tasks[0]));
        if (!tasks) ok = 0;
    }
    if (ok) {
        const uint32_t cache_budget =
            ds4_gpu_stream_expert_cache_configured_budget();
        uint32_t reserved_entries = g_stream_expert_cache_entry_count;

        for (uint32_t u = 0; u < unique_count; u++) {
            const uint32_t expert = (uint32_t)unique_ids[u];

            if ((uint64_t)expert > UINT64_MAX / gate_expert_bytes ||
                (uint64_t)expert > UINT64_MAX / down_expert_bytes) {
                fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n");
                ok = 0;
                break;
            }
            const uint64_t gate_rel = (uint64_t)expert * gate_expert_bytes;
            const uint64_t down_rel = (uint64_t)expert * down_expert_bytes;
            if (gate_rel > UINT64_MAX - gate_offset ||
                gate_rel > UINT64_MAX - up_offset ||
                down_rel > UINT64_MAX - down_offset) {
                fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n");
                ok = 0;
                break;
            }
            unique_gate_offsets[u] = gate_offset + gate_rel;
            unique_up_offsets[u] = up_offset + gate_rel;
            unique_down_offsets[u] = down_offset + down_rel;

            ds4_gpu_stream_expert_cache_entry *entry =
                ds4_gpu_stream_expert_cache_peek(model_map,
                                                 model_size,
                                                 layer,
                                                 expert,
                                                 n_total_expert,
                                                 n_selected,
                                                 unique_gate_offsets[u],
                                                 unique_up_offsets[u],
                                                 unique_down_offsets[u],
                                                 gate_expert_bytes,
                                                 down_expert_bytes);
            if (entry) {
                unique_entries[u] = entry;
                continue;
            }

            const int force_reuse =
                cache_budget != 0 && reserved_entries >= cache_budget;
            ds4_gpu_stream_expert_readahead_range(unique_gate_offsets[u],
                                                  gate_expert_bytes);
            ds4_gpu_stream_expert_readahead_range(unique_up_offsets[u],
                                                  gate_expert_bytes);
            ds4_gpu_stream_expert_readahead_range(unique_down_offsets[u],
                                                  down_expert_bytes);
            const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
            const int prepared =
                ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
                                                                 expert,
                                                                 layer,
                                                                 unique_ids,
                                                                 unique_count,
                                                                 gate_expert_bytes,
                                                                 down_expert_bytes,
                                                                 force_reuse,
                                                                 &gate_bufs[n_loads],
                                                                 &up_bufs[n_loads],
                                                                 &down_bufs[n_loads],
                                                                 &gate_inners[n_loads],
                                                                 &up_inners[n_loads],
                                                                 &down_inners[n_loads]);
            if (load_timing) {
                ds4_gpu_stream_expert_timing_note_prepare_buffer(
                        ds4_gpu_now_ms() - buffer_t0);
            }
            if (!prepared) {
                ok = 0;
                break;
            }
            if (!force_reuse && reserved_entries < UINT32_MAX) {
                reserved_entries++;
            }
            if (!gate_bufs[n_loads] ||
                !up_bufs[n_loads] ||
                !down_bufs[n_loads]) {
                ok = 0;
                break;
            }

            uint8_t *gate_dst = (uint8_t *)[gate_bufs[n_loads] contents] +
                                 gate_inners[n_loads];
            uint8_t *up_dst = (uint8_t *)[up_bufs[n_loads] contents] +
                               up_inners[n_loads];
            uint8_t *down_dst = (uint8_t *)[down_bufs[n_loads] contents] +
                                 down_inners[n_loads];
            if (!gate_dst || !up_dst || !down_dst) {
                ok = 0;
                break;
            }

            load_unique[n_loads] = u;
            const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
            tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
                .offset = unique_gate_offsets[u],
                .len = gate_expert_bytes,
                .dst = gate_dst,
            };
            tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
                .offset = unique_up_offsets[u],
                .len = gate_expert_bytes,
                .dst = up_dst,
            };
            tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
                .offset = unique_down_offsets[u],
                .len = down_expert_bytes,
                .dst = down_dst,
            };
            if (load_timing) {
                ds4_gpu_stream_expert_timing_note_prepare_task(
                        1,
                        ds4_gpu_now_ms() - task_t0);
            }
            n_loads++;
        }
    }
    if (ok && n_loads != 0) {
        if (load_timing_t0 != 0.0) {
            const double now_ms = ds4_gpu_now_ms();
            load_prepare_ms = now_ms - load_timing_t0;
            load_timing_t0 = now_ms;
        }
        uint64_t read_bytes = 0;
        double read_ms = 0.0;
        ok = ds4_gpu_stream_expert_pread_tasks(tasks,
                                               n_tasks,
                                               &read_bytes,
                                               &read_ms);
        if (ok) {
            ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms);
        }
        if (load_timing_t0 != 0.0) {
            load_timing_t0 = ds4_gpu_now_ms();
        }
        if (ok) {
            for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
                [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)];
                [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)];
                [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)];
            }
        }
        if (load_timing_t0 != 0.0) {
            const double now_ms = ds4_gpu_now_ms();
            load_modify_ms = now_ms - load_timing_t0;
            load_timing_t0 = now_ms;
        }
        if (ok && getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) {
            fprintf(stderr,
                    "ds4: Metal streaming batch expert parallel pread layer=%u experts=%u tensors=%u "
                    "threads=%u bytes=%.2f GiB wall=%.3f ms\n",
                    layer,
                    n_loads,
                    n_tasks,
                    ds4_gpu_stream_expert_pread_thread_count(n_tasks),
                    ds4_gpu_gib(read_bytes),
                    read_ms);
        }
        if (ok) {
            for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
                const uint32_t u = load_unique[load_i];
                const uint32_t expert = (uint32_t)unique_ids[u];
                ds4_gpu_stream_expert_cache_entry *entry =
                    ds4_gpu_stream_expert_cache_install_loaded(model_map,
                                                               model_size,
                                                               layer,
                                                               expert,
                                                               unique_gate_offsets[u],
                                                               unique_up_offsets[u],
                                                               unique_down_offsets[u],
                                                               gate_expert_bytes,
                                                               down_expert_bytes,
                                                               gate_bufs[load_i],
                                                               up_bufs[load_i],
                                                               down_bufs[load_i],
                                                               gate_inners[load_i],
                                                               up_inners[load_i],
                                                               down_inners[load_i]);
                if (!entry) {
                    ok = 0;
                    break;
                }
                unique_entries[u] = entry;
            }
        }
        if (load_timing_t0 != 0.0) {
            load_install_ms = ds4_gpu_now_ms() - load_timing_t0;
            ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms,
                                                          read_ms,
                                                          load_modify_ms,
                                                          load_install_ms);
        }
    }
    if (tasks) free(tasks);
    if (ok) {
        for (uint32_t u = 0; u < unique_count; u++) {
            ds4_gpu_stream_expert_cache_entry *entry = unique_entries[u];
            const uint32_t expert = (uint32_t)unique_ids[u];
            if (!entry) {
                const uint64_t gate_rel = unique_gate_offsets[u] - gate_offset;
                const uint64_t down_rel = unique_down_offsets[u] - down_offset;
                if (!*overflow_gate) {
                    const uint64_t gate_tensor_bytes =
                        (uint64_t)n_total_expert * gate_expert_bytes;
                    const uint64_t down_tensor_bytes =
                        (uint64_t)n_total_expert * down_expert_bytes;
                    uint64_t gate_view_inner = 0;
                    uint64_t up_view_inner = 0;
                    uint64_t down_view_inner = 0;
                    id<MTLBuffer> gv = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                gate_offset,
                                                                gate_tensor_bytes,
                                                                &gate_view_inner);
                    id<MTLBuffer> uv = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                up_offset,
                                                                gate_tensor_bytes,
                                                                &up_view_inner);
                    id<MTLBuffer> dv = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                down_offset,
                                                                down_tensor_bytes,
                                                                &down_view_inner);
                    if (!gv || !uv || !dv) {
                        fprintf(stderr,
                                "ds4: Metal streaming prefill batch selected addr "
                                "failed to map overflow expert views at layer %u\n",
                                layer);
                        ok = 0;
                        break;
                    }
                    *overflow_gate = gv;
                    *overflow_up = uv;
                    *overflow_down = dv;
                    overflow_gate_inner = gate_view_inner;
                    overflow_up_inner = up_view_inner;
                    overflow_down_inner = down_view_inner;
                }
                if (!ds4_gpu_stream_expert_cache_set_addr_slot(
                            layer,
                            expert,
                            *overflow_gate,
                            (NSUInteger)(overflow_gate_inner + gate_rel),
                            *overflow_up,
                            (NSUInteger)(overflow_up_inner + gate_rel),
                            *overflow_down,
                            (NSUInteger)(overflow_down_inner + down_rel))) {
                    ok = 0;
                    break;
                }
                view_served++;
                continue;
            }
            const uint32_t extra_uses =
                frequency[expert] > 0 ? frequency[expert] - 1u : 0;
            if (extra_uses != 0) {
                if (entry->use_count > UINT64_MAX - extra_uses) {
                    entry->use_count = UINT64_MAX;
                } else {
                    entry->use_count += extra_uses;
                }
                if (g_stream_expert_cache_hits > UINT64_MAX - extra_uses) {
                    g_stream_expert_cache_hits = UINT64_MAX;
                } else {
                    g_stream_expert_cache_hits += extra_uses;
                }
                if (g_stream_expert_cache_layer_hits[layer] >
                    UINT64_MAX - extra_uses) {
                    g_stream_expert_cache_layer_hits[layer] = UINT64_MAX;
                } else {
                    g_stream_expert_cache_layer_hits[layer] += extra_uses;
                }
            }
            resources[*n_resources] = entry;
            (*n_resources)++;
        }
    }
    free(ids);

    if (!ok || (*n_resources == 0 && view_served == 0)) {
        ds4_gpu_stream_expert_cache_clear_layer(layer);
        return 0;
    }
    if (!ds4_gpu_stream_expert_cache_addr_buffers(layer,
                                                  gate_addrs,
                                                  up_addrs,
                                                  down_addrs)) {
        ds4_gpu_stream_expert_cache_clear_layer(layer);
        return 0;
    }
    *unique_out = *n_resources + view_served;
    if (view_served != 0 &&
        ds4_gpu_stream_expert_timing_summary_enabled()) {
        fprintf(stderr,
                "ds4: Metal streaming prefill batch selected addr layer=%u "
                "served %u/%u unique experts via mapped views (cache budget %u)\n",
                layer,
                view_served,
                unique_count,
                ds4_gpu_stream_expert_cache_configured_budget());
    }
    if (profile) {
        const double t_done = ds4_gpu_now_ms();
        uint64_t per_expert_bytes = UINT64_MAX;
        uint64_t logical_bytes = UINT64_MAX;
        if (gate_expert_bytes <= (UINT64_MAX - down_expert_bytes) / 2ull) {
            per_expert_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
            if (per_expert_bytes == 0 ||
                *n_resources <= UINT64_MAX / per_expert_bytes) {
                logical_bytes = (uint64_t)(*n_resources) * per_expert_bytes;
            }
        }
        fprintf(stderr,
                "ds4: Metal streaming prefill batch selected addr layer=%u "
                "tokens=%u unique=%u read=%.3f ms wrap=%.3f ms bytes=%.2f GiB\n",
                layer,
                n_tokens,
                *n_resources,
                t_read - t0,
                t_done - t_read,
                ds4_gpu_gib(logical_bytes));
    }
    return 1;
}

int ds4_gpu_stream_expert_cache_seed_selected(
        const ds4_gpu_stream_expert_table *table,
        const int32_t                     *selected_ids,
        uint32_t                           n_selected) {
    if (!g_ssd_streaming_mode) return 1;
    if (!table) return 0;
    const void *model_map = table->model_map;
    const uint64_t model_size = table->model_size;
    const uint32_t layer = table->layer;
    const uint32_t n_total_expert = table->n_total_expert;
    const uint64_t gate_offset = table->gate_offset;
    const uint64_t up_offset = table->up_offset;
    const uint64_t down_offset = table->down_offset;
    const uint64_t gate_expert_bytes = table->gate_expert_bytes;
    const uint64_t down_expert_bytes = table->down_expert_bytes;
    if (!model_map || !selected_ids ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_selected == 0 ||
        n_selected > 6 ||
        n_total_expert == 0 ||
        n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes) ||
        ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                  n_total_expert,
                                                  n_selected) == 0) {
        return 1;
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;

    ds4_gpu_stream_expert_cache_note_selected_hotness(layer,
                                                      selected_ids,
                                                      n_selected);
    for (uint32_t i = 0; i < n_selected; i++) {
        if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
            fprintf(stderr,
                    "ds4: Metal prefill expert-cache seed selected expert id %d is outside 0..%u\n",
                    selected_ids[i],
                    n_total_expert);
            return 0;
        }
        const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i];
        if (expert_id > UINT64_MAX / gate_expert_bytes ||
            expert_id > UINT64_MAX / down_expert_bytes) {
            fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n");
            return 0;
        }
        const uint64_t gate_rel = expert_id * gate_expert_bytes;
        const uint64_t down_rel = expert_id * down_expert_bytes;
        if (gate_rel > UINT64_MAX - gate_offset ||
            gate_rel > UINT64_MAX - up_offset ||
            down_rel > UINT64_MAX - down_offset) {
            fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n");
            return 0;
        }

        if (!ds4_gpu_stream_expert_cache_get(model_map,
                                             model_size,
                                             layer,
                                             (uint32_t)selected_ids[i],
                                             n_total_expert,
                                             n_selected,
                                             gate_offset + gate_rel,
                                             up_offset + gate_rel,
                                             down_offset + down_rel,
                                             gate_expert_bytes,
                                             down_expert_bytes)) {
            return 0;
        }
    }
    ds4_gpu_stream_expert_cache_prune_layer(layer,
                                            n_total_expert,
                                            n_selected,
                                            selected_ids,
                                            n_selected);
    ds4_gpu_stream_expert_cache_prune_global(layer,
                                             selected_ids,
                                             n_selected);
    return 1;
}

int ds4_gpu_stream_expert_cache_seed_experts(
        const ds4_gpu_stream_expert_table *table,
        const int32_t                     *expert_ids,
        const uint32_t                    *expert_priorities,
        uint32_t                           n_experts) {
    if (!g_ssd_streaming_mode) return 1;
    if (!table) return 0;
    const void *model_map = table->model_map;
    const uint64_t model_size = table->model_size;
    const uint32_t layer = table->layer;
    const uint32_t n_total_expert = table->n_total_expert;
    const uint64_t gate_offset = table->gate_offset;
    const uint64_t up_offset = table->up_offset;
    const uint64_t down_offset = table->down_offset;
    const uint64_t gate_expert_bytes = table->gate_expert_bytes;
    const uint64_t down_expert_bytes = table->down_expert_bytes;
    if (!model_map || !expert_ids ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_experts == 0 ||
        n_total_expert == 0 ||
        n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
        !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                      down_expert_bytes) ||
        ds4_gpu_stream_expert_cache_effective_cap(layer,
                                                  n_total_expert,
                                                  1) == 0) {
        return 1;
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;

    ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
    uint32_t remaining = n_experts;
    while (remaining != 0) {
        const uint32_t batch =
            remaining > 6u ? 6u : remaining;
        remaining -= batch;

        int32_t selected_ids[6] = { -1, -1, -1, -1, -1, -1 };
        uint64_t gate_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        uint64_t up_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        uint64_t down_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        ds4_gpu_stream_expert_cache_entry *entries[6] = {
            NULL, NULL, NULL, NULL, NULL, NULL
        };
        uint32_t missing_mask = 0;

        for (uint32_t i = 0; i < batch; i++) {
            const int32_t expert = expert_ids[remaining + i];
            const uint32_t priority =
                expert_priorities ? expert_priorities[remaining + i] : 0;
            if (expert < 0 || (uint32_t)expert >= n_total_expert) {
                fprintf(stderr,
                        "ds4: Metal streaming hotlist seed expert id %d is outside 0..%u\n",
                        expert,
                        n_total_expert);
                return 0;
            }
            const uint64_t expert_id = (uint64_t)(uint32_t)expert;
            if (expert_id > UINT64_MAX / gate_expert_bytes ||
                expert_id > UINT64_MAX / down_expert_bytes) {
                fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n");
                return 0;
            }
            const uint64_t gate_rel = expert_id * gate_expert_bytes;
            const uint64_t down_rel = expert_id * down_expert_bytes;
            if (gate_rel > UINT64_MAX - gate_offset ||
                gate_rel > UINT64_MAX - up_offset ||
                down_rel > UINT64_MAX - down_offset) {
                fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n");
                return 0;
            }

            ds4_gpu_stream_expert_cache_note_route_hotness(
                    layer,
                    (uint32_t)expert,
                    priority != 0 ? priority : 1u);
            selected_ids[i] = expert;
            gate_abs_offsets[i] = gate_offset + gate_rel;
            up_abs_offsets[i] = up_offset + gate_rel;
            down_abs_offsets[i] = down_offset + down_rel;
            entries[i] = ds4_gpu_stream_expert_cache_peek(model_map,
                                                          model_size,
                                                          layer,
                                                          (uint32_t)expert,
                                                          n_total_expert,
                                                          batch,
                                                          gate_abs_offsets[i],
                                                          up_abs_offsets[i],
                                                          down_abs_offsets[i],
                                                          gate_expert_bytes,
                                                          down_expert_bytes);
            if (!entries[i]) missing_mask |= 1u << i;
            if (entries[i] && priority != 0 &&
                entries[i]->use_count < (uint64_t)priority) {
                entries[i]->use_count = (uint64_t)priority;
            }
        }

        if (missing_mask != 0 &&
            !ds4_gpu_stream_expert_cache_load_selected_missing(model_map,
                                                               model_size,
                                                               layer,
                                                               selected_ids,
                                                               n_total_expert,
                                                               batch,
                                                               gate_abs_offsets,
                                                               up_abs_offsets,
                                                               down_abs_offsets,
                                                               gate_expert_bytes,
                                                               down_expert_bytes,
                                                               missing_mask,
                                                               entries)) {
            return 0;
        }
        if (expert_priorities) {
            for (uint32_t i = 0; i < batch; i++) {
                const uint32_t priority = expert_priorities[remaining + i];
                if (entries[i] && priority != 0 &&
                    entries[i]->use_count < (uint64_t)priority) {
                    entries[i]->use_count = (uint64_t)priority;
                }
            }
        }
    }

    const uint32_t protect_n = n_experts < 6u ? n_experts : 6u;
    ds4_gpu_stream_expert_cache_prune_layer(layer,
                                            n_total_expert,
                                            1,
                                            expert_ids,
                                            protect_n);
    ds4_gpu_stream_expert_cache_prune_global(layer,
                                             expert_ids,
                                             protect_n);
    return 1;
}

static uint32_t ds4_gpu_q4_expert_table_group_size(uint32_t n_total_expert) {
    const char *env = getenv("DS4_METAL_Q4_EXPERT_TABLE_GROUP_SIZE");
    if (!env || !env[0]) return 1;
    char *end = NULL;
    unsigned long v = strtoul(env, &end, 10);
    if (end == env || *end != '\0' || v < 2 || v > n_total_expert) {
        return 1;
    }
    return (uint32_t)v;
}

static bool ds4_gpu_q4_table_queue_residency_requested(void) {
    return getenv("DS4_METAL_Q4_TABLE_QUEUE_RESIDENCY_SET") != NULL;
}

static bool ds4_gpu_q4_table_queue_residency_available(void) {
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        return g_queue && [g_queue respondsToSelector:@selector(addResidencySet:)];
    }
#endif
    return false;
}

static bool ds4_gpu_q4_non_streaming_opt_in_enabled(void) {
    return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL ||
           getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL ||
           getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
           getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL ||
           getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL ||
           getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL;
}

static bool ds4_gpu_q4_selected_paths_allowed(void) {
    if (g_ssd_streaming_mode) return true;
    if (g_glm_model_mode) return false;
    return ds4_gpu_q4_non_streaming_opt_in_enabled();
}

int ds4_gpu_pro_q4_expert_table_auto_available(void) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    return (g_ssd_streaming_mode ||
            getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL) &&
           getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL &&
           getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL &&
           ds4_gpu_q4_table_queue_residency_available();
}

static bool ds4_gpu_q4_table_queue_residency_enabled(bool auto_queue_residency) {
    return (auto_queue_residency || ds4_gpu_q4_table_queue_residency_requested()) &&
           ds4_gpu_q4_table_queue_residency_available();
}

static bool ds4_gpu_q4_table_model_residency_enabled(void) {
    return getenv("DS4_METAL_Q4_TABLE_MODEL_RESIDENCY_SET") != NULL;
}

static bool ds4_gpu_pro_q4_expert_indirect_shape_supported(
        uint32_t n_total_expert,
        uint32_t n_expert,
        uint64_t gate_tensor_bytes,
        uint64_t down_tensor_bytes) {
    const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull;
    return n_total_expert == 384 &&
           n_expert == 6 &&
           gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
           down_tensor_bytes >= q4_selected_min_tensor_bytes;
}

static bool ds4_gpu_pro_q4_expert_table_auto_enabled(
        uint32_t n_total_expert,
        uint32_t n_expert,
        uint64_t gate_tensor_bytes,
        uint64_t down_tensor_bytes) {
    /*
     * This path lets the shader choose among exact per-expert resources.
     * It is only automatic when a Metal queue residency set can make every
     * indirect expert resource visible up front; otherwise the selected
     * exact-slice path remains the fallback.
     */
    return ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert,
                                                          n_expert,
                                                          gate_tensor_bytes,
                                                          down_tensor_bytes) &&
           ds4_gpu_pro_q4_expert_table_auto_available();
}

static bool ds4_gpu_pro_q4_expert_address_auto_enabled(
        uint32_t n_total_expert,
        uint32_t n_expert,
        uint64_t gate_tensor_bytes,
        uint64_t down_tensor_bytes) {
    /*
     * GPU-address expert tables are useful for experiments, but they are not
     * safe as an automatic path unless every indirect expert resource is made
     * visible to Metal. Keep this behind an explicit opt-in while the selected
     * active-slice path remains the correctness/performance baseline.
     */
    return getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL &&
           ds4_gpu_q4_selected_paths_allowed() &&
           ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert,
                                                          n_expert,
                                                          gate_tensor_bytes,
                                                          down_tensor_bytes) &&
           g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil &&
           g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil &&
           getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL &&
           getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") == NULL &&
           getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL &&
           ds4_gpu_q4_table_queue_residency_available();
}

static bool ds4_gpu_q4_table_bind_anchors_enabled(void) {
    return getenv("DS4_METAL_Q4_TABLE_BIND_ANCHORS") != NULL;
}

static int ds4_gpu_use_model_residency_set(id<MTLCommandBuffer> cb) {
    if (!ds4_gpu_q4_table_model_residency_enabled()) return 1;
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        if (cb && g_model_residency_set && [cb respondsToSelector:@selector(useResidencySet:)]) {
            [cb useResidencySet:g_model_residency_set];
            return 1;
        }
    }
#endif
    fprintf(stderr, "ds4: Metal Q4 table model residency set is not available\n");
    return 0;
}

static int ds4_gpu_bind_q4_expert_table_anchors(
        id<MTLComputeCommandEncoder> enc,
        DS4MetalQ4ExpertTable       *table,
        NSUInteger                   first_index,
        NSUInteger                   max_count) {
    if (!ds4_gpu_q4_table_bind_anchors_enabled()) return 1;
    if (!enc || !table || !table.expertBuffers) return 0;

    const NSUInteger count = [table.expertBuffers count];
    if (count > max_count) {
        fprintf(stderr,
                "ds4: Metal Q4 table anchor count %lu exceeds available slots %lu\n",
                (unsigned long)count,
                (unsigned long)max_count);
        return 0;
    }
    for (NSUInteger i = 0; i < count; i++) {
        [enc setBuffer:[table.expertBuffers objectAtIndex:i]
                offset:0
               atIndex:first_index + i];
    }
    return 1;
}

static id ds4_gpu_q4_expert_table_residency_set(NSMutableArray<id<MTLBuffer>> *buffers) {
    if (!buffers || [buffers count] == 0 ||
        getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL ||
        getenv("DS4_METAL_Q4_TABLE_PER_TENSOR_RESIDENCY_SET") == NULL) {
        return nil;
    }
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init];
        desc.label = @"ds4_q4_expert_table";
        desc.initialCapacity = [buffers count];
        NSError *error = nil;
        id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error];
        if (!residency_set) {
            fprintf(stderr, "ds4: Metal Q4 expert table residency set creation failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            return nil;
        }
        for (id<MTLBuffer> buffer in buffers) {
            [residency_set addAllocation:buffer];
        }
        [residency_set commit];
        [residency_set requestResidency];
        return residency_set;
    }
#endif
    return nil;
}

static void ds4_gpu_q4_residency_add_table(id residency_set,
                                           DS4MetalQ4ExpertTable *table) {
#if TARGET_OS_OSX
    if (!residency_set || !table) return;
    if (@available(macOS 15.0, *)) {
        if (table.argumentBuffer) {
            [residency_set addAllocation:table.argumentBuffer];
        }
        if (table.addressBuffer) {
            [residency_set addAllocation:table.addressBuffer];
        }
        for (id<MTLBuffer> buffer in table.expertBuffers) {
            [residency_set addAllocation:buffer];
        }
    }
#else
    (void)residency_set;
    (void)table;
#endif
}

static id ds4_gpu_q4_expert_layer_residency_set(DS4MetalQ4ExpertTable *gate_table,
                                                DS4MetalQ4ExpertTable *up_table,
                                                DS4MetalQ4ExpertTable *down_table,
                                                bool auto_queue_residency) {
    const bool queue_residency =
        ds4_gpu_q4_table_queue_residency_enabled(auto_queue_residency);
    if (!gate_table || !up_table || !down_table ||
        !g_device || !g_q4_expert_layer_residency_cache ||
        (getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL &&
         !queue_residency)) {
        return nil;
    }
#if TARGET_OS_OSX
    if (@available(macOS 15.0, *)) {
        NSString *key = [NSString stringWithFormat:@"%p:%p:%p",
                         gate_table, up_table, down_table];
        DS4MetalQ4LayerResidency *cached =
            [g_q4_expert_layer_residency_cache objectForKey:key];
        if (cached) return cached.residencySet;

        const NSUInteger capacity =
            [gate_table.expertBuffers count] +
            [up_table.expertBuffers count] +
            [down_table.expertBuffers count] + 3u;
        MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init];
        desc.label = @"ds4_q4_expert_layer";
        desc.initialCapacity = capacity;
        NSError *error = nil;
        id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error];
        if (!residency_set) {
            fprintf(stderr,
                    "ds4: Metal Q4 expert layer residency set creation failed: %s\n",
                    [[error localizedDescription] UTF8String]);
            return nil;
        }

        ds4_gpu_q4_residency_add_table(residency_set, gate_table);
        ds4_gpu_q4_residency_add_table(residency_set, up_table);
        ds4_gpu_q4_residency_add_table(residency_set, down_table);
        [residency_set commit];
        [residency_set requestResidency];

        DS4MetalQ4LayerResidency *entry = [DS4MetalQ4LayerResidency new];
        entry.residencySet = residency_set;
        if (queue_residency) {
            [g_queue addResidencySet:residency_set];
            entry.addedToQueue = YES;
        }
        [g_q4_expert_layer_residency_cache setObject:entry forKey:key];
        return residency_set;
    }
#endif
    return nil;
}

static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table(
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                tensor_offset,
        uint64_t                expert_bytes,
        uint32_t                n_total_expert,
        id<MTLArgumentEncoder>  encoder);

static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    tensor_offset,
        uint64_t    expert_bytes,
        uint32_t    n_total_expert);

int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size,
                                     uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset,
                                     uint64_t gate_expert_bytes, uint64_t down_expert_bytes,
                                     uint32_t n_total_expert) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!model_map || model_size == 0 || gate_expert_bytes == 0 ||
        down_expert_bytes == 0 || n_total_expert == 0) {
        return 0;
    }
    if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
        fprintf(stderr, "ds4: Metal Q4 expert table preload byte size overflow\n");
        return 0;
    }

    const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes;
    const bool address_auto =
        ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert,
                                                   6,
                                                   gate_tensor_bytes,
                                                   down_tensor_bytes);
    const bool table_auto =
        ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
                                                 6,
                                                 gate_tensor_bytes,
                                                 down_tensor_bytes);
    if (!address_auto && !table_auto) {
        return 1;
    }

    if (address_auto) {
        @autoreleasepool {
            DS4MetalQ4ExpertTable *gate_table =
                ds4_gpu_q4_expert_address_table(model_map,
                                                model_size,
                                                gate_offset,
                                                gate_expert_bytes,
                                                n_total_expert);
            DS4MetalQ4ExpertTable *up_table =
                ds4_gpu_q4_expert_address_table(model_map,
                                                model_size,
                                                up_offset,
                                                gate_expert_bytes,
                                                n_total_expert);
            DS4MetalQ4ExpertTable *down_table =
                ds4_gpu_q4_expert_address_table(model_map,
                                                model_size,
                                                down_offset,
                                                down_expert_bytes,
                                                n_total_expert);
            if (!gate_table || !up_table || !down_table) {
                return 0;
            }
            id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table,
                                                                 up_table,
                                                                 down_table,
                                                                 true);
            if (!residency) {
                fprintf(stderr, "ds4: Metal Q4 expert address table preload failed to create queue residency set\n");
                return 0;
            }
        }
        return 1;
    }

    if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder ||
        !g_moe_table_q4_sum_down_encoder) {
        fprintf(stderr, "ds4: Metal Q4 expert table preload missing argument encoders\n");
        return 0;
    }

    @autoreleasepool {
        DS4MetalQ4ExpertTable *gate_table =
            ds4_gpu_q4_expert_table(model_map,
                                    model_size,
                                    gate_offset,
                                    gate_expert_bytes,
                                    n_total_expert,
                                    g_moe_table_q4_pair_gate_encoder);
        DS4MetalQ4ExpertTable *up_table =
            ds4_gpu_q4_expert_table(model_map,
                                    model_size,
                                    up_offset,
                                    gate_expert_bytes,
                                    n_total_expert,
                                    g_moe_table_q4_pair_up_encoder);
        DS4MetalQ4ExpertTable *down_table =
            ds4_gpu_q4_expert_table(model_map,
                                    model_size,
                                    down_offset,
                                    down_expert_bytes,
                                    n_total_expert,
                                    g_moe_table_q4_sum_down_encoder);
        if (!gate_table || !up_table || !down_table) {
            return 0;
        }
        id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table,
                                                             up_table,
                                                             down_table,
                                                             true);
        if (!residency) {
            fprintf(stderr, "ds4: Metal Q4 expert table preload failed to create queue residency set\n");
            return 0;
        }
    }
    return 1;
}

static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table(
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                tensor_offset,
        uint64_t                expert_bytes,
        uint32_t                n_total_expert,
        id<MTLArgumentEncoder>  encoder) {
    if (!model_map || !g_device || !g_q4_expert_table_cache || !encoder ||
        model_size == 0 || expert_bytes == 0 ||
        n_total_expert == 0 || n_total_expert > 384) {
        return nil;
    }
    if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) {
        fprintf(stderr, "ds4: Metal Q4 expert table byte size overflow\n");
        return nil;
    }
    const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes;
    if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) {
        fprintf(stderr, "ds4: Metal Q4 expert table is outside the mapped model\n");
        return nil;
    }

    const uint32_t table_group_size =
        ds4_gpu_q4_expert_table_group_size(n_total_expert);
    NSString *key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu:%u:%llu:%u",
                     model_map,
                     (unsigned long long)model_size,
                     (unsigned long long)tensor_offset,
                     (unsigned long long)expert_bytes,
                     n_total_expert,
                     (unsigned long long)[encoder encodedLength],
                     table_group_size];
    DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key];
    if (cached) return cached;

    id<MTLBuffer> arg_buffer =
        [g_device newBufferWithLength:[encoder encodedLength]
                               options:MTLResourceStorageModeShared];
    if (!arg_buffer) {
        fprintf(stderr, "ds4: Metal Q4 expert table argument buffer allocation failed\n");
        return nil;
    }
    arg_buffer.label = @"ds4_q4_expert_table";
    [encoder setArgumentBuffer:arg_buffer offset:0];

    NSMutableArray<id<MTLBuffer>> *expert_buffers =
        [NSMutableArray arrayWithCapacity:table_group_size > 1 ?
            (n_total_expert + table_group_size - 1u) / table_group_size :
            n_total_expert];
    if (!expert_buffers) return nil;

    if (table_group_size > 1) {
        for (uint32_t first = 0; first < n_total_expert; first += table_group_size) {
            const uint32_t remaining = n_total_expert - first;
            const uint32_t group_n =
                remaining < table_group_size ? remaining : table_group_size;
            if ((uint64_t)first > UINT64_MAX / expert_bytes ||
                (uint64_t)group_n > UINT64_MAX / expert_bytes) {
                fprintf(stderr, "ds4: Metal Q4 expert table group byte overflow\n");
                return nil;
            }
            const uint64_t rel = (uint64_t)first * expert_bytes;
            const uint64_t group_bytes = (uint64_t)group_n * expert_bytes;
            if (rel > UINT64_MAX - tensor_offset) {
                fprintf(stderr, "ds4: Metal Q4 expert table group offset overflow\n");
                return nil;
            }
            uint64_t inner = 0;
            id<MTLBuffer> group_buf =
                ds4_gpu_wrap_model_range(model_map,
                                         model_size,
                                         tensor_offset + rel,
                                         group_bytes,
                                         &inner);
            if (!group_buf) return nil;
            for (uint32_t j = 0; j < group_n; j++) {
                const uint64_t expert_inner = inner + (uint64_t)j * expert_bytes;
                [encoder setBuffer:group_buf offset:(NSUInteger)expert_inner atIndex:first + j];
            }
            [expert_buffers addObject:group_buf];
        }
    } else {
        for (uint32_t i = 0; i < n_total_expert; i++) {
            const uint64_t rel = (uint64_t)i * expert_bytes;
            if (rel > UINT64_MAX - tensor_offset) {
                fprintf(stderr, "ds4: Metal Q4 expert table offset overflow\n");
                return nil;
            }
            uint64_t inner = 0;
            id<MTLBuffer> expert_buf =
                ds4_gpu_wrap_model_exact_range_owned(model_map,
                                                     model_size,
                                                     tensor_offset + rel,
                                                     expert_bytes,
                                                     &inner);
            if (!expert_buf) return nil;
            [encoder setBuffer:expert_buf offset:(NSUInteger)inner atIndex:i];
            [expert_buffers addObject:expert_buf];
        }
    }
    [arg_buffer didModifyRange:NSMakeRange(0, [encoder encodedLength])];

    if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) {
        fprintf(stderr,
                "ds4: Metal Q4 expert table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n",
                n_total_expert,
                table_group_size,
                (unsigned long)[expert_buffers count],
                ds4_gpu_mib(expert_bytes));
    }

    DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new];
    table.argumentBuffer = arg_buffer;
    table.expertBuffers = expert_buffers;
    table.residencySet = ds4_gpu_q4_expert_table_residency_set(expert_buffers);
    table.nExpert = n_total_expert;
    table.expertBytes = expert_bytes;
    [g_q4_expert_table_cache setObject:table forKey:key];
    return table;
}

static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table(
        const void *model_map,
        uint64_t    model_size,
        uint64_t    tensor_offset,
        uint64_t    expert_bytes,
        uint32_t    n_total_expert) {
    if (!model_map || !g_device || !g_q4_expert_table_cache ||
        model_size == 0 || expert_bytes == 0 ||
        n_total_expert == 0 || n_total_expert > 384) {
        return nil;
    }
    if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) {
        fprintf(stderr, "ds4: Metal Q4 expert address table byte size overflow\n");
        return nil;
    }
    const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes;
    if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) {
        fprintf(stderr, "ds4: Metal Q4 expert address table is outside the mapped model\n");
        return nil;
    }

    const uint32_t table_group_size =
        ds4_gpu_q4_expert_table_group_size(n_total_expert);
    NSString *key = [NSString stringWithFormat:@"addr:%p:%llu:%llu:%llu:%u:%u",
                     model_map,
                     (unsigned long long)model_size,
                     (unsigned long long)tensor_offset,
                     (unsigned long long)expert_bytes,
                     n_total_expert,
                     table_group_size];
    DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key];
    if (cached) return cached;

    id<MTLBuffer> address_buffer =
        [g_device newBufferWithLength:(NSUInteger)n_total_expert * sizeof(uint64_t)
                               options:MTLResourceStorageModeShared];
    if (!address_buffer) {
        fprintf(stderr, "ds4: Metal Q4 expert address table allocation failed\n");
        return nil;
    }
    address_buffer.label = @"ds4_q4_expert_address_table";
    uint64_t *addresses = (uint64_t *)[address_buffer contents];
    if (!addresses) return nil;

    NSMutableArray<id<MTLBuffer>> *expert_buffers =
        [NSMutableArray arrayWithCapacity:table_group_size > 1 ?
            (n_total_expert + table_group_size - 1u) / table_group_size :
            n_total_expert];
    if (!expert_buffers) return nil;

#if TARGET_OS_OSX
    if (@available(macOS 13.0, *)) {
        for (uint32_t first = 0; first < n_total_expert; first += table_group_size) {
            const uint32_t remaining = n_total_expert - first;
            const uint32_t group_n =
                remaining < table_group_size ? remaining : table_group_size;
            if ((uint64_t)first > UINT64_MAX / expert_bytes ||
                (uint64_t)group_n > UINT64_MAX / expert_bytes) {
                fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n");
                return nil;
            }
            const uint64_t rel = (uint64_t)first * expert_bytes;
            const uint64_t group_bytes = (uint64_t)group_n * expert_bytes;
            if (rel > UINT64_MAX - tensor_offset) {
                fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n");
                return nil;
            }
            uint64_t inner = 0;
            id<MTLBuffer> group_buf = nil;
            if (table_group_size > 1) {
                group_buf = ds4_gpu_wrap_model_range(model_map,
                                                     model_size,
                                                     tensor_offset + rel,
                                                     group_bytes,
                                                     &inner);
            } else {
                group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map,
                                                                 model_size,
                                                                 tensor_offset + rel,
                                                                 expert_bytes,
                                                                 &inner);
            }
            if (!group_buf) return nil;
            const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner;
            for (uint32_t j = 0; j < group_n; j++) {
                addresses[first + j] = base_address + (uint64_t)j * expert_bytes;
            }
            [expert_buffers addObject:group_buf];
        }
    } else
#endif
    {
        fprintf(stderr, "ds4: Metal GPU addresses require macOS 13 or newer\n");
        return nil;
    }

    [address_buffer didModifyRange:NSMakeRange(0,
        (NSUInteger)n_total_expert * sizeof(uint64_t))];

    if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) {
        fprintf(stderr,
                "ds4: Metal Q4 expert address table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n",
                n_total_expert,
                table_group_size,
                (unsigned long)[expert_buffers count],
                ds4_gpu_mib(expert_bytes));
    }

    DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new];
    table.addressBuffer = address_buffer;
    table.expertBuffers = expert_buffers;
    table.nExpert = n_total_expert;
    table.expertBytes = expert_bytes;
    [g_q4_expert_table_cache setObject:table forKey:key];
    return table;
}

static void ds4_gpu_use_q4_expert_table_resources(
        id<MTLCommandBuffer>        cb,
        id<MTLComputeCommandEncoder> enc,
        DS4MetalQ4ExpertTable       *table,
        bool                         queue_residency) {
    if (!enc || !table) return;
    if (table.argumentBuffer) {
        [enc useResource:table.argumentBuffer usage:MTLResourceUsageRead];
    }
    if (table.addressBuffer) {
        [enc useResource:table.addressBuffer usage:MTLResourceUsageRead];
    }
    if (!queue_residency &&
        getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL &&
        table.residencySet &&
        cb &&
        [cb respondsToSelector:@selector(useResidencySet:)]) {
        [cb useResidencySet:table.residencySet];
    }
    if (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") == NULL &&
        getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") == NULL) {
        return;
    }
    const NSUInteger count = [table.expertBuffers count];
    for (NSUInteger base = 0; base < count; base += 64u) {
        const NSUInteger n = count - base < 64u ? count - base : 64u;
        __unsafe_unretained id<MTLResource> resources[64];
        for (NSUInteger i = 0; i < n; i++) {
            resources[i] = [table.expertBuffers objectAtIndex:base + i];
        }
        [enc useResources:resources count:n usage:MTLResourceUsageRead];
    }
}

int ds4_gpu_indexer_score_one_tensor(
        ds4_gpu_tensor       *scores,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *weights,
        const ds4_gpu_tensor *index_comp,
        uint32_t                n_comp,
        uint32_t                n_head,
        uint32_t                head_dim,
        float                   scale) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!scores || !q || !weights || !index_comp ||
        n_comp == 0 || n_head == 0 || head_dim == 0) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
        const uint64_t weight_bytes = (uint64_t)n_head * sizeof(float);
        const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
        const uint64_t score_bytes = (uint64_t)n_comp * sizeof(float);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(index_comp);
        id<MTLBuffer> scorebuf = ds4_gpu_tensor_buffer(scores);
        if (!qbuf || !wbuf || !compbuf || !scorebuf ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(weights) < weight_bytes ||
            ds4_gpu_tensor_bytes(index_comp) < comp_bytes ||
            ds4_gpu_tensor_bytes(scores) < score_bytes) {
            fprintf(stderr, "ds4: Metal graph indexer score received undersized buffers\n");
            return 0;
        }

        if (n_head == 64 && head_dim == 128) {
            id<MTLComputePipelineState> direct_pipeline =
                ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_direct_pipeline,
                                        "kernel_dsv4_indexer_score_one_direct");
            if (!direct_pipeline) return 0;

            ds4_gpu_dsv4_indexer_scores_fused_args args = {
                .n_comp = n_comp,
                .n_tokens = 1,
                .n_head = n_head,
                .head_dim = head_dim,
                .pos0 = 0,
                .ratio = 4,
                .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float),
                .q_head_stride = (uint64_t)head_dim * sizeof(float),
                .weights_token_stride = (uint64_t)n_head * sizeof(float),
                .index_row_stride = (uint64_t)head_dim * sizeof(float),
                .score_token_stride = (uint64_t)n_comp * sizeof(float),
                .scale = scale,
            };

            int owned = 0;
            id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;
            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:direct_pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
            [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
            [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3];
            [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
            [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(n_comp, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer direct score")) return 0;
            return 1;
        }

        const uint64_t head_score_bytes = (uint64_t)n_comp * n_head * sizeof(float);
        if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_head_scores_buffer,
                                             &g_indexer_head_scores_bytes,
                                             (NSUInteger)head_score_bytes,
                                             "ds4_indexer_head_scores")) {
            return 0;
        }

        ds4_gpu_q8_0_matvec_args dot_args =
            ds4_gpu_make_f32_mv_args(head_dim, n_comp, n_head);
        ds4_gpu_mv_dispatch dot_dispatch =
            ds4_gpu_make_plain_mv_dispatch(head_dim, 1);
        dot_args.nr0 = dot_dispatch.nr0;
        id<MTLComputePipelineState> dot_pipeline =
            ds4_gpu_get_mul_mv_pipeline(dot_dispatch.function_name, dot_dispatch.nsg);
        if (!dot_pipeline) return 0;
        ds4_gpu_dsv4_indexer_weighted_sum_args sum_args = {
            .ne00 = (int64_t)n_comp,
            .ne01 = 1,
            .ne02 = (int64_t)n_head,
            .nb00 = sizeof(float),
            .nb01 = (uint64_t)n_comp * sizeof(float),
            .nb02 = (uint64_t)n_comp * sizeof(float),
            .ne10 = (int64_t)n_head,
            .ne11 = 1,
            .nb10 = sizeof(float),
            .nb11 = (uint64_t)n_head * sizeof(float),
            .ne0 = (int64_t)n_comp,
            .ne1 = 1,
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_comp * sizeof(float),
            .scale = scale,
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:dot_pipeline];
        [enc setBytes:&dot_args length:sizeof(dot_args) atIndex:0];
        [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:1];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2];
        [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:3];
        if (dot_dispatch.smem) {
            [enc setThreadgroupMemoryLength:dot_dispatch.smem atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + (NSUInteger)dot_dispatch.nr0 - 1u) / (NSUInteger)dot_dispatch.nr0,
                                              n_head,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dot_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_indexer_weighted_sum_pipeline];
        [enc setBytes:&sum_args length:sizeof(sum_args) atIndex:0];
        [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:1];
        [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
        [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 255u) / 256u, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer score")) return 0;
    }

    return 1;
}

static int ds4_gpu_indexer_scores_batch_tensor(
        ds4_gpu_tensor       *scores,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *weights,
        const ds4_gpu_tensor *index_comp,
        uint32_t                n_comp,
        uint32_t                n_tokens,
        uint32_t                pos0,
        uint32_t                n_head,
        uint32_t                head_dim,
        uint32_t                ratio,
        float                   scale) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!scores || !q || !weights || !index_comp ||
        n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || ratio == 0) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
        const uint64_t weight_bytes = (uint64_t)n_tokens * n_head * sizeof(float);
        const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
        const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(index_comp);
        id<MTLBuffer> scorebuf = ds4_gpu_tensor_buffer(scores);
        if (!qbuf || !wbuf || !compbuf || !scorebuf ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(weights) < weight_bytes ||
            ds4_gpu_tensor_bytes(index_comp) < comp_bytes ||
            ds4_gpu_tensor_bytes(scores) < score_bytes) {
            fprintf(stderr, "ds4: Metal graph indexer prefill scores received undersized buffers\n");
            return 0;
        }
        if (head_dim != 128) {
            fprintf(stderr, "ds4: Metal fused DS4 indexer scores expect 128-wide rows\n");
            return 0;
        }
        /*
         * The NAX/TensorOps score builder is a prefill-only win.  At small
         * batches and in one-token decode the setup cost is not amortized, so
         * those paths keep the older direct/tiled score kernels.
         */
        const bool use_nax = ds4_gpu_mpp_available() && n_tokens >= 16u;
        id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
            use_nax ? "kernel_dsv4_indexer_scores_nax" :
            (g_quality_mode ? "kernel_dsv4_indexer_scores_tiled_f32"
                            : "kernel_dsv4_indexer_scores_tiled"));
        if (!pipeline) return 0;

        ds4_gpu_dsv4_indexer_scores_fused_args args = {
            .n_comp = n_comp,
            .n_tokens = n_tokens,
            .n_head = n_head,
            .head_dim = head_dim,
            .pos0 = pos0,
            .ratio = ratio,
            .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float),
            .q_head_stride = (uint64_t)head_dim * sizeof(float),
            .weights_token_stride = (uint64_t)n_head * sizeof(float),
            .index_row_stride = (uint64_t)head_dim * sizeof(float),
            .score_token_stride = (uint64_t)n_comp * sizeof(float),
            .scale = scale,
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
        [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3];
        [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
        if (use_nax) {
            const NSUInteger q_shared = 2u * 32u * 32u;
            const NSUInteger k_shared = 32u * 128u;
            const NSUInteger dot_shared = 32u * 32u;
            [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) +
                                            dot_shared * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u,
                                                  ((NSUInteger)n_tokens + 15u) / 16u,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        } else if (g_quality_mode) {
            const NSUInteger q_shared = 8u * 128u;
            const NSUInteger k_shared = 32u * 128u;
            const NSUInteger dot_shared = 8u * 32u;
            [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u,
                                                  ((NSUInteger)n_tokens + 7u) / 8u,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
        } else {
            const NSUInteger q_shared = 8u * 128u;
            const NSUInteger k_shared = 32u * 128u;
            const NSUInteger dot_shared = 8u * 32u;
            [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) +
                                            dot_shared * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u,
                                                  ((NSUInteger)n_tokens + 7u) / 8u,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer prefill scores")) return 0;
    }

    return 1;
}

int ds4_gpu_indexer_scores_prefill_tensor(
        ds4_gpu_tensor       *scores,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *weights,
        const ds4_gpu_tensor *index_comp,
        uint32_t                n_comp,
        uint32_t                n_tokens,
        uint32_t                n_head,
        uint32_t                head_dim,
        uint32_t                ratio,
        float                   scale) {
    return ds4_gpu_indexer_scores_batch_tensor(scores,
                                                 q,
                                                 weights,
                                                 index_comp,
                                                 n_comp,
                                                 n_tokens,
                                                 0,
                                                 n_head,
                                                 head_dim,
                                                 ratio,
                                                 scale);
}

int ds4_gpu_indexer_scores_decode_batch_tensor(
        ds4_gpu_tensor       *scores,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *weights,
        const ds4_gpu_tensor *index_comp,
        uint32_t                n_comp,
        uint32_t                n_tokens,
        uint32_t                pos0,
        uint32_t                n_head,
        uint32_t                head_dim,
        uint32_t                ratio,
        float                   scale) {
    return ds4_gpu_indexer_scores_batch_tensor(scores,
                                                 q,
                                                 weights,
                                                 index_comp,
                                                 n_comp,
                                                 n_tokens,
                                                 pos0,
                                                 n_head,
                                                 head_dim,
                                                 ratio,
                                                 scale);
}

int ds4_gpu_indexer_topk_tensor(
        ds4_gpu_tensor       *selected,
        const ds4_gpu_tensor *scores,
        uint32_t                n_comp,
        uint32_t                n_tokens,
        uint32_t                top_k) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || top_k > n_comp) return 0;

    @autoreleasepool {
        const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float);
        const uint64_t selected_bytes = (uint64_t)top_k * n_tokens * sizeof(uint32_t);
        id<MTLBuffer> scorebuf = ds4_gpu_tensor_buffer(scores);
        id<MTLBuffer> selbuf = ds4_gpu_tensor_buffer(selected);
        if (!scorebuf || !selbuf ||
            ds4_gpu_tensor_bytes(scores) < score_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n");
            return 0;
        }
        NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup;
        if (max_threads == 0) max_threads = 256;
        int32_t nth = 1;
        while ((uint32_t)nth < n_comp && (uint64_t)2u * (uint64_t)nth <= (uint64_t)max_threads) {
            nth *= 2;
        }
        const int32_t npr = (int32_t)((n_comp + (uint32_t)nth - 1u) / (uint32_t)nth);
        const int32_t block_top_k = (int32_t)(top_k < (uint32_t)nth ? top_k : (uint32_t)nth);
        int32_t work_width = (int32_t)top_k;
        if (npr > 1) {
            const int32_t last_block = (int32_t)n_comp - (npr - 1) * nth;
            work_width = (npr - 1) * block_top_k + (last_block < block_top_k ? last_block : block_top_k);
        }
        const uint64_t scratch_row_bytes = (uint64_t)work_width * sizeof(uint32_t);
        const bool one_pass = npr <= 1;
        const uint64_t scratch_bytes = one_pass ? scratch_row_bytes * n_tokens :
            2u * scratch_row_bytes * n_tokens;
        if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_topk_buffer,
                                             &g_indexer_topk_bytes,
                                             (NSUInteger)scratch_bytes,
                                             "ds4_indexer_topk")) {
            return 0;
        }

        ds4_gpu_kargs_argsort args = {
            .ne00 = (int32_t)n_comp,
            .ne01 = (int32_t)n_tokens,
            .ne02 = 1,
            .ne03 = 1,
            .nb00 = sizeof(float),
            .nb01 = (uint64_t)n_comp * sizeof(float),
            .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float),
            .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float),
            .ne0 = work_width,
            .ne1 = (int32_t)n_tokens,
            .ne2 = 1,
            .ne3 = 1,
            .top_k = block_top_k,
        };
        // kernel_argsort_f32_i32_desc stages the block's scores behind the
        // index array: nth int32 indices + nth float scores.
        const NSUInteger smem = (((NSUInteger)nth * (sizeof(int32_t) + sizeof(float))) + 15u) & ~(NSUInteger)15u;

        NSUInteger cur_off = 0;
        NSUInteger next_off = (NSUInteger)scratch_row_bytes * n_tokens;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_argsort_f32_i32_desc_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1];
        [enc setBuffer:one_pass ? selbuf : g_indexer_topk_buffer
              offset:one_pass ? ds4_gpu_tensor_offset(selected) : cur_off
             atIndex:2];
        [enc setThreadgroupMemoryLength:smem atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)npr * n_tokens, 1, 1)
             threadsPerThreadgroup:MTLSizeMake((NSUInteger)nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        int32_t len = block_top_k;
        while (len < work_width) {
            const int32_t nm = (work_width + 2 * len - 1) / (2 * len);
            const bool final_merge = nm == 1;
            NSUInteger merge_threads = g_argsort_merge_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup;
            if (merge_threads == 0 || merge_threads > 512u) merge_threads = 512u;
            if (merge_threads > (NSUInteger)len) merge_threads = (NSUInteger)len;
            if (merge_threads == 0) merge_threads = 1;

            ds4_gpu_kargs_argsort_merge merge_args = {
                .ne00 = (int64_t)n_comp,
                .ne01 = (int64_t)n_tokens,
                .ne02 = 1,
                .ne03 = 1,
                .nb00 = sizeof(float),
                .nb01 = (uint64_t)n_comp * sizeof(float),
                .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float),
                .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float),
                .ne0 = work_width,
                .ne1 = (int32_t)n_tokens,
                .ne2 = 1,
                .ne3 = 1,
                .top_k = nm == 1 ? (int32_t)top_k : work_width,
                .len = len,
            };

            enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:g_argsort_merge_f32_i32_desc_pipeline];
            [enc setBytes:&merge_args length:sizeof(merge_args) atIndex:0];
            [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1];
            [enc setBuffer:g_indexer_topk_buffer offset:cur_off atIndex:2];
            [enc setBuffer:final_merge ? selbuf : g_indexer_topk_buffer
                  offset:final_merge ? ds4_gpu_tensor_offset(selected) : next_off
                 atIndex:3];
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nm * n_tokens, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(merge_threads, 1, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            const NSUInteger tmp = cur_off;
            cur_off = next_off;
            next_off = tmp;
            len <<= 1;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k")) return 0;
    }

    return 1;
}

int ds4_gpu_argmax_tensor(
        ds4_gpu_tensor       *out_idx,
        const ds4_gpu_tensor *logits,
        uint32_t                n_vocab) {
    if (!out_idx || !logits || n_vocab == 0) return 0;
    if (ds4_gpu_tensor_bytes(out_idx) < sizeof(int32_t) ||
        ds4_gpu_tensor_bytes(logits) < (uint64_t)n_vocab * sizeof(float)) {
        fprintf(stderr, "ds4: Metal graph argmax received undersized buffers\n");
        return 0;
    }

    return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1, 1);
}

int ds4_gpu_dsv4_topk_mask_tensor(
        ds4_gpu_tensor       *mask,
        const ds4_gpu_tensor *topk,
        uint32_t                n_comp,
        uint32_t                n_tokens,
        uint32_t                top_k) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0) return 0;

    @autoreleasepool {
        const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t);
        const uint64_t mask_bytes = (uint64_t)n_comp * n_tokens * sizeof(float);
        id<MTLBuffer> topkbuf = ds4_gpu_tensor_buffer(topk);
        id<MTLBuffer> maskbuf = ds4_gpu_tensor_buffer(mask);
        if (!topkbuf || !maskbuf ||
            ds4_gpu_tensor_bytes(topk) < topk_bytes ||
            ds4_gpu_tensor_bytes(mask) < mask_bytes) {
            fprintf(stderr, "ds4: Metal dsv4 top-k mask received undersized buffers\n");
            return 0;
        }

        ds4_gpu_dsv4_topk_mask_args args = {
            .ne00 = (int64_t)top_k,
            .ne01 = (int64_t)n_tokens,
            .nb00 = sizeof(int32_t),
            .nb01 = (uint64_t)top_k * sizeof(int32_t),
            .ne0 = (int64_t)n_comp,
            .ne1 = (int64_t)n_tokens,
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_comp * sizeof(float),
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_topk_mask_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1];
        [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)n_comp * n_tokens) + 255u) / 256u, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_topk_mask_scatter_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1];
        [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)top_k * n_tokens) + 255u) / 256u, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "dsv4 top-k mask")) return 0;
    }

    return 1;
}

static int ds4_gpu_matmul_q8_0_legacy_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok,
        bool                    prefer_decode_mpp,
        bool                    force_model_view) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if ((in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
        const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal Q8_0 tensor matmul received undersized activation buffers\n");
            return 0;
        }

        const uint64_t blocks = in_dim / 32;
        const uint64_t row_bytes = blocks * 34;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal Q8_0 tensor matmul range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = force_model_view ?
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset) :
            ds4_gpu_wrap_q8_decode_model_range(model_map,
                                               model_size,
                                               weight_offset,
                                               weight_bytes,
                                               n_tok,
                                               &inner_offset);
        if (!wbuf) {
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (n_tok == 1) {
            if (ds4_gpu_mpp_available() &&
                (prefer_decode_mpp || getenv("DS4_METAL_Q8_DECODE_MPP") != NULL) &&
                getenv("DS4_METAL_DISABLE_Q8_DECODE_MPP") == NULL &&
                (in_dim % 64u) == 0) {
                const char *nax_fn = "kernel_mul_mm_q8_0_f32_nax_direct_rhs";
                id<MTLComputePipelineState> mpp_pipeline =
                    ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false);
                if (mpp_pipeline) {
                    ds4_gpu_mul_mm_args args =
                        ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

                    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                    [enc setComputePipelineState:mpp_pipeline];
                    [enc setBytes:&args length:sizeof(args) atIndex:0];
                    [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
                    [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
                    [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                    [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0];
                    [enc dispatchThreadgroups:MTLSizeMake(1u,
                                                          ((NSUInteger)out_dim + 63u) / 64u,
                                                          1u)
                         threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
                    ds4_gpu_end_compute_encoder(cb, enc);

                    if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 decode MPP matmul")) {
                        return 0;
                    }
                    return 1;
                }
                ds4_gpu_warn_mpp_fallback();
            }

            ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
            ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
            if (out_dim > 65536u) mv_dispatch.nsg = 8;
            const bool force_output_nr4 =
                getenv("DS4_METAL_ENABLE_OUTPUT_Q8_NR4") != NULL;
            const bool output_shape =
                in_dim == 4096u && out_dim == 129280u;
            const bool use_output_nr4 =
                !g_quality_mode && (out_dim & 3u) == 0u &&
                (force_output_nr4 ||
                 (output_shape && ds4_gpu_device_name_contains("M3"))) &&
                getenv("DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4") == NULL;
            if (use_output_nr4) {
                mv_dispatch.function_name =
                    "kernel_mul_mv_q8_0_f32_nr4";
                mv_dispatch.nr0 = 4;
                mv_dispatch.smem = 32u * 4u * sizeof(float);
            }
            mv_args.nr0 = mv_dispatch.nr0;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
            if (!pipeline) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
                                                  1,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matvec")) {
                return 0;
            }
            return 1;
        }

        const uint64_t mv_ext_max_tokens =
            ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u);
        if (n_tok <= mv_ext_max_tokens && (in_dim % 128u) == 0) {
            const int16_t nsg = 2;
            const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
            const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
            const char *fn_name = ds4_gpu_mv_ext_name(1, r1ptg);
            id<MTLComputePipelineState> pipeline =
                fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
            if (!pipeline) return 0;

            const int16_t nypsg = 32 / nxpsg;
            const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
            ds4_gpu_mul_mv_ext_args args =
                ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes);

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
                                                  ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor mul_mv_ext")) {
                return 0;
            }
            return 1;
        }

        /*
         * Dense Q8_0 prefill is the cleanest DS4 TensorOps shape: M/N/K are
         * aligned and the RHS activation matrix is already dense.  The retained
         * kernel dequantizes each 64x32 weight tile to half in threadgroup
         * memory, then uses direct-RHS MPP for the activation tile.  This avoids
         * staging RHS into threadgroup memory and was the direct replacement for
         * the slower generic MPP prototype.
         */
        if (ds4_gpu_mpp_available() &&
            n_tok >= 32u &&
            (in_dim % 64u) == 0 &&
            (out_dim % 64u) == 0 &&
            (n_tok % 32u) == 0) {
            uint64_t nax_tile_n = 32u;
            if ((n_tok % 128u) == 0) {
                nax_tile_n = 128u;
            } else if ((n_tok % 64u) == 0) {
                nax_tile_n = 64u;
            }
            const char *nax_fn = nax_tile_n == 128u
                ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128"
                : (nax_tile_n == 64u
                    ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64"
                    : "kernel_mul_mm_q8_0_f32_nax_direct_rhs");
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false);
            if (pipeline) {
                ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

                id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
                [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
                [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0];
                [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n),
                                                      (NSUInteger)out_dim / 64u,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);

                if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) {
                    return 0;
                }
                return 1;
            }
            ds4_gpu_warn_mpp_fallback();
        }

        const bool bc_inp = (in_dim % 32u) != 0;
        const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_q8_0_f32", bc_inp, bc_out);
        if (!pipeline) return 0;

        ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u,
                                              ((NSUInteger)out_dim + 63u) / 64u,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matmul")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_matmul_q8_0_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if ((in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
        return 0;
    }

    const int profile_requested =
        n_tok > 8u && ds4_gpu_env_bool("DS4_METAL_Q8_PREFILL_PROFILE") > 0;
    int profile_prefill = 0;
    int split_batch_for_profile = 0;
    const char *profile_label = NULL;
    char profile_label_buf[128];
    char profile_fallback[128];
    if (profile_requested) {
        snprintf(profile_fallback, sizeof(profile_fallback),
                 "q8 weight_off=%llu in=%llu out=%llu tok=%llu",
                 (unsigned long long)weight_offset,
                 (unsigned long long)in_dim,
                 (unsigned long long)out_dim,
                 (unsigned long long)n_tok);
        snprintf(profile_label_buf, sizeof(profile_label_buf), "%s", profile_fallback);
        profile_label = profile_label_buf;
        const char *profile_filter = getenv("DS4_METAL_Q8_PREFILL_PROFILE_FILTER");
        profile_prefill =
            profile_requested &&
            (!profile_filter || !profile_filter[0] ||
             strstr(profile_label, profile_filter) != NULL);
    }
    if (profile_prefill) {
        if (g_batch_cb) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
            split_batch_for_profile = 1;
        }
    }

    const double profile_t0 = profile_prefill ? ds4_gpu_now_ms() : 0.0;
    int ok = ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size,
                                               weight_offset, in_dim, out_dim,
                                               x, n_tok, false, false);
    if (profile_prefill) {
        if (split_batch_for_profile && ds4_gpu_end_commands() == 0) {
            ok = 0;
        }
        const double elapsed_ms = ds4_gpu_now_ms() - profile_t0;
        fprintf(stderr,
                "ds4: Metal Q8_0 prefill profile %s in=%llu out=%llu tok=%llu %.3f ms\n",
                profile_label ? profile_label : profile_fallback,
                (unsigned long long)in_dim,
                (unsigned long long)out_dim,
                (unsigned long long)n_tok,
                elapsed_ms);
        if (split_batch_for_profile && ds4_gpu_begin_commands() == 0) {
            ok = 0;
        }
    }
    return ok;
}

int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(
        ds4_gpu_tensor       *out,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint64_t              in_dim,
        uint64_t              out_dim,
        const ds4_gpu_tensor *x,
        uint32_t              n_rows) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !x || !model_map || n_rows == 0 ||
        n_rows > INT32_MAX || in_dim == 0 || out_dim == 0 ||
        (in_dim & 31u) != 0 || in_dim > UINT32_MAX ||
        out_dim > UINT32_MAX ||
        in_dim > UINT64_MAX / n_rows / sizeof(float) ||
        out_dim > UINT64_MAX / n_rows / sizeof(float) ||
        ds4_gpu_tensor_bytes(x) <
            (uint64_t)n_rows * in_dim * sizeof(float) ||
        ds4_gpu_tensor_bytes(out) <
            (uint64_t)n_rows * out_dim * sizeof(float)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t blocks = in_dim / 32u;
        const uint64_t row_bytes = blocks * 34u;
        if (!xbuf || !outbuf ||
            out_dim > UINT64_MAX / row_bytes) {
            return 0;
        }
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (weight_offset > model_size ||
            weight_bytes > model_size - weight_offset) {
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = ds4_gpu_wrap_q8_decode_model_range(
                model_map, model_size, weight_offset, weight_bytes, 1u,
                &inner_offset);
        if (!wbuf) return 0;

        ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        if (out_dim > 65536u) dispatch.nsg = 8;
        ds4_gpu_q8_0_matvec_args args =
            ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
        args.ne11 = (int32_t)n_rows;
        args.nb12 = (uint64_t)n_rows * in_dim * sizeof(float);
        args.nb13 = args.nb12;
        args.ne1 = (int32_t)n_rows;
        args.nr0 = dispatch.nr0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline(dispatch.function_name, dispatch.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:dispatch.smem atIndex:0];
        [enc dispatchThreadgroups:
                MTLSizeMake(((NSUInteger)out_dim +
                             (NSUInteger)dispatch.nr0 - 1u) /
                                (NSUInteger)dispatch.nr0,
                            (NSUInteger)n_rows,
                            1)
             threadsPerThreadgroup:
                MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        return ds4_gpu_finish_command_buffer(
                cb, owned, "Q8_0 exact decode-row matvec");
    }
}

int ds4_gpu_matmul_q8_0_decode_mpp_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size,
                                             weight_offset, in_dim, out_dim,
                                             x, n_tok, true, false);
}

int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size,
                                             weight_offset, in_dim, out_dim,
                                             x, n_tok, true, true);
}

static const char *ds4_gpu_q4_mv_ext_name(uint32_t weight_type, int16_t r1ptg) {
    const char *prefix = NULL;
    if (weight_type == DS4_METAL_TENSOR_Q4_K) {
        prefix = "kernel_mul_mv_ext_q4_K_f32_r1_";
    } else if (weight_type == DS4_METAL_TENSOR_Q4_0) {
        prefix = "kernel_mul_mv_ext_q4_0_f32_r1_";
    } else {
        return NULL;
    }
    switch (r1ptg) {
    case 1: return weight_type == DS4_METAL_TENSOR_Q4_K ?
        "kernel_mul_mv_ext_q4_K_f32_r1_1" : "kernel_mul_mv_ext_q4_0_f32_r1_1";
    case 2: return weight_type == DS4_METAL_TENSOR_Q4_K ?
        "kernel_mul_mv_ext_q4_K_f32_r1_2" : "kernel_mul_mv_ext_q4_0_f32_r1_2";
    case 3: return weight_type == DS4_METAL_TENSOR_Q4_K ?
        "kernel_mul_mv_ext_q4_K_f32_r1_3" : "kernel_mul_mv_ext_q4_0_f32_r1_3";
    case 4: return weight_type == DS4_METAL_TENSOR_Q4_K ?
        "kernel_mul_mv_ext_q4_K_f32_r1_4" : "kernel_mul_mv_ext_q4_0_f32_r1_4";
    case 5: return weight_type == DS4_METAL_TENSOR_Q4_K ?
        "kernel_mul_mv_ext_q4_K_f32_r1_5" : "kernel_mul_mv_ext_q4_0_f32_r1_5";
    default:
        (void)prefix;
        return NULL;
    }
}

static const char *ds4_gpu_q4_mm_name(uint32_t weight_type) {
    switch (weight_type) {
    case DS4_METAL_TENSOR_Q4_0: return "kernel_mul_mm_q4_0_f32";
    case DS4_METAL_TENSOR_Q4_K: return "kernel_mul_mm_q4_K_f32";
    default: return NULL;
    }
}

static const char *ds4_gpu_q4_nax_name(uint32_t weight_type, uint64_t tile_n) {
    if (weight_type == DS4_METAL_TENSOR_Q4_0) {
        return tile_n == 128u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128" :
               tile_n == 64u  ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64" :
                                 "kernel_mul_mm_q4_0_f32_nax_direct_rhs";
    }
    if (weight_type == DS4_METAL_TENSOR_Q4_K) {
        return tile_n == 128u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128" :
               tile_n == 64u  ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64" :
                                 "kernel_mul_mm_q4_K_f32_nax_direct_rhs";
    }
    return NULL;
}

static int ds4_gpu_matmul_quant_impl_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                weight_type,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok,
        bool                    prefer_decode_mpp,
        bool                    force_model_view) {
    if (weight_type == DS4_METAL_TENSOR_Q8_0) {
        return ds4_gpu_matmul_q8_0_legacy_tensor(out,
                                                 model_map,
                                                 model_size,
                                                 weight_offset,
                                                 in_dim,
                                                 out_dim,
                                                 x,
                                                 n_tok,
                                                 prefer_decode_mpp,
                                                 force_model_view);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !x || !model_map ||
        in_dim == 0 || out_dim == 0 || n_tok == 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
        return 0;
    }

    uint64_t row_bytes = 0;
    if (!ds4_gpu_quant_row_bytes(weight_type, (uint32_t)in_dim, &row_bytes)) {
        fprintf(stderr, "ds4: Metal quant matmul received unsupported type/dim (%u, in=%llu)\n",
                weight_type,
                (unsigned long long)in_dim);
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
        const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal quant tensor matmul received undersized activation buffers\n");
            return 0;
        }

        if (out_dim > UINT64_MAX / row_bytes) return 0;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal quant tensor matmul range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = force_model_view ?
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        /*
         * Small-batch Q4_K goes to the classic (llama.cpp-style) matvec:
         * the mul_mv_ext family tops out around 220 GB/s on M5 for the GLM
         * DenseQ4 decode shapes while this impl streams 530-650 GB/s
         * (misc/q4mv_bench.m). Falls through to ext when unavailable.
         */
        if (weight_type == DS4_METAL_TENSOR_Q4_K &&
            n_tok <= 8u &&
            (in_dim % 256u) == 0 &&
            getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) {
            const int16_t nsg = 2;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8);
            if (pipeline) {
                ds4_gpu_q8_0_matvec_args args = {
                    .ne00 = (int32_t)in_dim,
                    .ne01 = (int32_t)out_dim,
                    .ne02 = 1,
                    .nb00 = 1,
                    .nb01 = row_bytes,
                    .nb02 = row_bytes * out_dim,
                    .nb03 = row_bytes * out_dim,
                    .ne10 = (int32_t)in_dim,
                    .ne11 = (int32_t)n_tok,
                    .ne12 = 1,
                    .nb10 = sizeof(float),
                    .nb11 = in_dim * sizeof(float),
                    .nb12 = in_dim * n_tok * sizeof(float),
                    .nb13 = in_dim * n_tok * sizeof(float),
                    .ne0 = (int32_t)out_dim,
                    .ne1 = (int32_t)n_tok,
                    .nr0 = 2,
                    .r2 = 1,
                    .r3 = 1,
                };
                const uint64_t rows_ptg = (uint64_t)nsg * 2u;

                id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
                [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
                [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                [enc setThreadgroupMemoryLength:32 atIndex:0];
                [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg,
                                                      (NSUInteger)n_tok,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);

                if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K classic mul_mv")) return 0;
                return 1;
            }
        }

        if (n_tok <= 8u && (in_dim % 128u) == 0) {
            const int16_t nsg = 2;
            const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
            const int16_t r1ptg = (n_tok == 1u) ? 1 : ds4_gpu_mv_ext_r1ptg(n_tok);
            const char *fn_name = ds4_gpu_q4_mv_ext_name(weight_type, r1ptg);
            id<MTLComputePipelineState> pipeline =
                fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
            if (pipeline) {
                const int16_t nypsg = 32 / nxpsg;
                const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
                ds4_gpu_mul_mv_ext_args args =
                    ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, row_bytes, row_bytes);

                id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
                [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
                [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
                                                      ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);

                if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor mul_mv_ext")) return 0;
                return 1;
            }
        }

        if (ds4_gpu_mpp_available() &&
            n_tok >= 32u &&
            (in_dim % 64u) == 0 &&
            (out_dim % 64u) == 0 &&
            (n_tok % 32u) == 0) {
            uint64_t nax_tile_n = 32u;
            if ((n_tok % 128u) == 0) {
                nax_tile_n = 128u;
            } else if ((n_tok % 64u) == 0) {
                nax_tile_n = 64u;
            }
            const char *nax_fn = ds4_gpu_q4_nax_name(weight_type, nax_tile_n);
            id<MTLComputePipelineState> pipeline =
                nax_fn ? ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false) : nil;
            if (pipeline) {
                ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

                id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
                [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
                [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0];
                [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n),
                                                      (NSUInteger)out_dim / 64u,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);

                if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 NAX tensor matmul")) return 0;
                return 1;
            }
            ds4_gpu_warn_mpp_fallback();
        }

        const char *mm_fn = ds4_gpu_q4_mm_name(weight_type);
        const bool bc_inp = (in_dim % 32u) != 0;
        const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0;
        id<MTLComputePipelineState> pipeline =
            mm_fn ? ds4_gpu_get_mul_mm_pipeline(mm_fn, bc_inp, bc_out) : nil;
        if (!pipeline) return 0;

        ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u,
                                              ((NSUInteger)out_dim + 63u) / 64u,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor matmul")) return 0;
    }

    return 1;
}

int ds4_gpu_matmul_quant_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                weight_type,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    return ds4_gpu_matmul_quant_impl_tensor(out,
                                            model_map,
                                            model_size,
                                            weight_offset,
                                            weight_type,
                                            in_dim,
                                            out_dim,
                                            x,
                                            n_tok,
                                            false,
                                            false);
}

int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                weight_type,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    /* GLM decode dense matvecs: the classic Q8_0 kernels stream 570-613
     * GB/s on the GLM shapes while the MPP/nax matrix path measures ~150
     * GB/s at n_tok=1 (TP decode 8.75 -> 16.64 t/s on the IQ2+Q8 gguf).
     * MPP stays available as an opt-in via DS4_METAL_Q8_DECODE_MPP. */
    const bool prefer_mpp = getenv("DS4_METAL_Q8_DECODE_MPP") != NULL;
    return ds4_gpu_matmul_quant_impl_tensor(out,
                                            model_map,
                                            model_size,
                                            weight_offset,
                                            weight_type,
                                            in_dim,
                                            out_dim,
                                            x,
                                            n_tok,
                                            prefer_mpp,
                                            true);
}

int ds4_gpu_matmul_quant_rows_scalar_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                weight_type,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    return ds4_gpu_matmul_quant_tensor(out,
                                       model_map,
                                       model_size,
                                       weight_offset,
                                       weight_type,
                                       in_dim,
                                       out_dim,
                                       x,
                                       n_tok);
}

int ds4_gpu_matmul_q8_0_rows_scalar_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    if (n_tok == 1) {
        return ds4_gpu_matmul_q8_0_tensor(out,
                                          model_map,
                                          model_size,
                                          weight_offset,
                                          in_dim,
                                          out_dim,
                                          x,
                                          1);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !x || !model_map ||
        n_tok == 0 || n_tok > INT32_MAX ||
        (in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX ||
        in_dim > UINT64_MAX / sizeof(float) ||
        out_dim > UINT64_MAX / sizeof(float)) {
        return 0;
    }

    const uint64_t x_row_bytes = in_dim * sizeof(float);
    const uint64_t out_row_bytes = out_dim * sizeof(float);
    if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) ||
        (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t x_bytes = n_tok * x_row_bytes;
        const uint64_t out_bytes = n_tok * out_row_bytes;
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul received undersized activation buffers\n");
            return 0;
        }

        const uint64_t blocks = in_dim / 32;
        if (blocks > UINT64_MAX / 34u) return 0;
        const uint64_t row_bytes = blocks * 34u;
        if (out_dim > UINT64_MAX / row_bytes) return 0;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner_offset);
        if (!wbuf) return 0;

        ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        if (out_dim > 65536u) mv_dispatch.nsg = 8;
        mv_args.nr0 = mv_dispatch.nr0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
        [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];

        const NSUInteger x_base = ds4_gpu_tensor_offset(x);
        const NSUInteger out_base = ds4_gpu_tensor_offset(out);
        for (uint64_t t = 0; t < n_tok; t++) {
            [enc setBuffer:xbuf offset:x_base + (NSUInteger)(t * x_row_bytes) atIndex:2];
            [enc setBuffer:outbuf offset:out_base + (NSUInteger)(t * out_row_bytes) atIndex:3];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
                                                      (NSUInteger)mv_dispatch.nr0,
                                                  1,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 scalar-row matvecs")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_matmul_q8_0_pair_tensor(
        ds4_gpu_tensor       *out0,
        ds4_gpu_tensor       *out1,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight0_offset,
        uint64_t                weight1_offset,
        uint64_t                in_dim,
        uint64_t                out0_dim,
        uint64_t                out1_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out0 || !out1 || !model_map || !x || n_tok != 1 ||
        out0_dim == 0 || out1_dim == 0 || (in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> out0buf = ds4_gpu_tensor_buffer(out0);
        id<MTLBuffer> out1buf = ds4_gpu_tensor_buffer(out1);
        const uint64_t x_bytes = in_dim * sizeof(float);
        const uint64_t out0_bytes = out0_dim * sizeof(float);
        const uint64_t out1_bytes = out1_dim * sizeof(float);
        if (!xbuf || !out0buf || !out1buf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out0) < out0_bytes ||
            ds4_gpu_tensor_bytes(out1) < out1_bytes) {
            fprintf(stderr, "ds4: Metal paired Q8_0 matvec received undersized activation buffers\n");
            return 0;
        }

        const uint64_t row_bytes = (in_dim / 32u) * 34u;
        const uint64_t weight0_bytes = out0_dim * row_bytes;
        const uint64_t weight1_bytes = out1_dim * row_bytes;
        if (weight0_offset > model_size || weight0_bytes > model_size - weight0_offset ||
            weight1_offset > model_size || weight1_bytes > model_size - weight1_offset) {
            fprintf(stderr, "ds4: Metal paired Q8_0 matvec range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner0 = 0;
        uint64_t inner1 = 0;
        id<MTLBuffer> weight0buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                     weight0_offset, weight0_bytes, &inner0);
        id<MTLBuffer> weight1buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                     weight1_offset, weight1_bytes, &inner1);
        if (!weight0buf || !weight1buf) return 0;

        ds4_gpu_mv_dispatch dispatch0 = ds4_gpu_make_q8_0_mv_dispatch();
        ds4_gpu_mv_dispatch dispatch1 = ds4_gpu_make_q8_0_mv_dispatch();
        if (out0_dim > 65536u) dispatch0.nsg = 8;
        if (out1_dim > 65536u) dispatch1.nsg = 8;
        /* A common threadgroup shape is required to retain each standalone
         * reduction tree. Mixed 4/8-simdgroup extents use the existing fallback. */
        if (dispatch0.nsg != dispatch1.nsg) return 0;

        ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, out0_dim);
        ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, out1_dim);
        args0.nr0 = dispatch0.nr0;
        args1.nr0 = dispatch1.nr0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_pair", dispatch0.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args0 length:sizeof(args0) atIndex:0];
        [enc setBytes:&args1 length:sizeof(args1) atIndex:1];
        [enc setBuffer:weight0buf offset:(NSUInteger)inner0 atIndex:2];
        [enc setBuffer:weight1buf offset:(NSUInteger)inner1 atIndex:3];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
        [enc setBuffer:out0buf offset:ds4_gpu_tensor_offset(out0) atIndex:5];
        [enc setBuffer:out1buf offset:ds4_gpu_tensor_offset(out1) atIndex:6];
        [enc setThreadgroupMemoryLength:2u * dispatch0.smem atIndex:0];
        const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim;
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)max_out_dim +
                                               (NSUInteger)dispatch0.nr0 - 1u) /
                                              (NSUInteger)dispatch0.nr0,
                                              1,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch0.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "paired Q8_0 matvec")) return 0;
    }

    return 1;
}

int ds4_gpu_matmul_q8_0_f16_out_tensor(
        ds4_gpu_tensor       *out_h,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    (void)out_h; (void)model_map; (void)model_size; (void)weight_offset;
    (void)in_dim; (void)out_dim; (void)x; (void)n_tok;
    return 0;
}

static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl(
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        float                   clamp,
        int                     store_gate_up,
        bool                    force_model_view) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!mid || !x || !model_map ||
        (store_gate_up && (!gate || !up)) ||
        (in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX ||
        !isfinite(clamp) || clamp < 0.0f) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> gatebuf = store_gate_up ?
            ds4_gpu_tensor_buffer(gate) : midbuf;
        id<MTLBuffer> upbuf = store_gate_up ?
            ds4_gpu_tensor_buffer(up) : midbuf;
        const uint64_t x_bytes = in_dim * sizeof(float);
        const uint64_t out_bytes = out_dim * sizeof(float);
        if (!xbuf || !gatebuf || !upbuf || !midbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            (store_gate_up && ds4_gpu_tensor_bytes(gate) < out_bytes) ||
            (store_gate_up && ds4_gpu_tensor_bytes(up) < out_bytes) ||
            ds4_gpu_tensor_bytes(mid) < out_bytes) {
            fprintf(stderr, "ds4: Metal shared expert fused gate/up received undersized activation buffers\n");
            return 0;
        }

        const uint64_t blocks = in_dim / 32;
        const uint64_t row_bytes = blocks * 34;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (gate_offset > model_size || weight_bytes > model_size - gate_offset ||
            up_offset > model_size || weight_bytes > model_size - up_offset) {
            fprintf(stderr, "ds4: Metal shared expert fused gate/up range is outside the mapped model\n");
            return 0;
        }

        const bool exact_decode_views =
            !force_model_view &&
            getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL &&
            getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL;
        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        id<MTLBuffer> gate_wbuf = exact_decode_views ?
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           gate_offset,
                                           weight_bytes,
                                           &gate_inner) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     gate_offset,
                                     weight_bytes,
                                     &gate_inner);
        id<MTLBuffer> up_wbuf = exact_decode_views ?
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           up_offset,
                                           weight_bytes,
                                           &up_inner) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     up_offset,
                                     weight_bytes,
                                     &up_inner);
        if (!gate_wbuf || !up_wbuf) return 0;

        ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        args.nr0 = mv_dispatch.nr0;
        const char *fn_name = store_gate_up ?
            (mv_dispatch.nr0 >= 4 ?
             "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" :
             "kernel_dsv4_shared_gate_up_swiglu_q8_0") :
            (mv_dispatch.nr0 >= 4 ?
             "kernel_dsv4_shared_mid_swiglu_q8_0_r4" :
             "kernel_dsv4_shared_mid_swiglu_q8_0");
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1];
        [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:gatebuf offset:(store_gate_up ?
                                       ds4_gpu_tensor_offset(gate) :
                                       ds4_gpu_tensor_offset(mid)) atIndex:4];
        [enc setBuffer:upbuf offset:(store_gate_up ?
                                     ds4_gpu_tensor_offset(up) :
                                     ds4_gpu_tensor_offset(mid)) atIndex:5];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
        [enc setBytes:&clamp length:sizeof(clamp) atIndex:7];
        [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
                                                  (NSUInteger)mv_dispatch.nr0,
                                              1,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb,
                                           owned,
                                           store_gate_up ?
                                           "shared expert fused gate/up" :
                                           "shared expert fused mid")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        float                   clamp) {
    return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate,
                                                   up,
                                                   mid,
                                                   model_map,
                                                   model_size,
                                                   gate_offset,
                                                   up_offset,
                                                   in_dim,
                                                   out_dim,
                                                   x,
                                                   clamp,
                                                   1,
                                                   false);
}

int ds4_gpu_shared_mid_swiglu_q8_0_tensor(
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        float                   clamp) {
    return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(NULL,
                                                   NULL,
                                                   mid,
                                                   model_map,
                                                   model_size,
                                                   gate_offset,
                                                   up_offset,
                                                   in_dim,
                                                   out_dim,
                                                   x,
                                                   clamp,
                                                   0,
                                                   true);
}

int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        float                   clamp) {
    return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate,
                                                   up,
                                                   mid,
                                                   model_map,
                                                   model_size,
                                                   gate_offset,
                                                   up_offset,
                                                   in_dim,
                                                   out_dim,
                                                   x,
                                                   clamp,
                                                   1,
                                                   true);
}

int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok,
        float                   clamp) {
    if (n_tok == 1) {
        return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate,
                                                         up,
                                                         mid,
                                                         model_map,
                                                         model_size,
                                                         gate_offset,
                                                         up_offset,
                                                         in_dim,
                                                         out_dim,
                                                         x,
                                                         clamp);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!gate || !up || !mid || !x || !model_map ||
        n_tok == 0 ||
        (in_dim & 31u) != 0 || (in_dim % 128u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX ||
        !isfinite(clamp) || clamp < 0.0f) {
        return 0;
    }

    const uint64_t mv_ext_max_tokens =
        ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u);
    if (n_tok > mv_ext_max_tokens) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
        id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
        const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
        if (!xbuf || !gatebuf || !upbuf || !midbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(gate) < out_bytes ||
            ds4_gpu_tensor_bytes(up) < out_bytes ||
            ds4_gpu_tensor_bytes(mid) < out_bytes) {
            fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows received undersized activation buffers\n");
            return 0;
        }

        const uint64_t blocks = in_dim / 32;
        const uint64_t row_bytes = blocks * 34;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (gate_offset > model_size || weight_bytes > model_size - gate_offset ||
            up_offset > model_size || weight_bytes > model_size - up_offset) {
            fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows range is outside the mapped model\n");
            return 0;
        }

        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        id<MTLBuffer> gate_wbuf =
            ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner);
        id<MTLBuffer> up_wbuf =
            ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner);
        if (!gate_wbuf || !up_wbuf) return 0;

        const int16_t nsg = 2;
        const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
        const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
        const char *fn_name = ds4_gpu_mv_ext_q8_pair_swiglu_name(r1ptg);
        id<MTLComputePipelineState> pipeline =
            fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
        if (!pipeline) return 0;

        const int16_t nypsg = 32 / nxpsg;
        const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
        ds4_gpu_mul_mv_ext_args args =
            ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes);

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1];
        [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4];
        [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
        [enc setBytes:&clamp length:sizeof(clamp) atIndex:7];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) /
                                                  (NSUInteger)r0ptg,
                                              ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) /
                                                  (NSUInteger)r1ptg,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "fused Q8_0 gate/up rows")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor(
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok,
        float                   clamp) {
    if (n_tok == 1) {
        return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate,
                                                         up,
                                                         mid,
                                                         model_map,
                                                         model_size,
                                                         gate_offset,
                                                         up_offset,
                                                         in_dim,
                                                         out_dim,
                                                         x,
                                                         clamp);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!gate || !up || !mid || !x || !model_map ||
        n_tok == 0 ||
        (in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX ||
        in_dim > UINT64_MAX / sizeof(float) ||
        out_dim > UINT64_MAX / sizeof(float) ||
        !isfinite(clamp) || clamp < 0.0f) {
        return 0;
    }

    const uint64_t x_row_bytes = in_dim * sizeof(float);
    const uint64_t out_row_bytes = out_dim * sizeof(float);
    if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) ||
        (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
        id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        const uint64_t x_bytes = n_tok * x_row_bytes;
        const uint64_t out_bytes = n_tok * out_row_bytes;
        if (!xbuf || !gatebuf || !upbuf || !midbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(gate) < out_bytes ||
            ds4_gpu_tensor_bytes(up) < out_bytes ||
            ds4_gpu_tensor_bytes(mid) < out_bytes) {
            fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up received undersized activation buffers\n");
            return 0;
        }

        const uint64_t blocks = in_dim / 32;
        if (blocks > UINT64_MAX / 34u) return 0;
        const uint64_t row_bytes = blocks * 34u;
        if (out_dim > UINT64_MAX / row_bytes) return 0;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (gate_offset > model_size || weight_bytes > model_size - gate_offset ||
            up_offset > model_size || weight_bytes > model_size - up_offset) {
            fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up range is outside the mapped model\n");
            return 0;
        }

        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        id<MTLBuffer> gate_wbuf =
            ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner);
        id<MTLBuffer> up_wbuf =
            ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner);
        if (!gate_wbuf || !up_wbuf) return 0;

        ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        args.nr0 = mv_dispatch.nr0;
        args.ne11 = (int32_t)n_tok;
        args.nb12 = n_tok * x_row_bytes;
        args.nb13 = args.nb12;
        args.ne1 = (int32_t)n_tok;
        const char *fn_name = mv_dispatch.nr0 >= 4 ?
            "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" :
            "kernel_dsv4_shared_gate_up_swiglu_q8_0";
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1];
        [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4];
        [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
        [enc setBytes:&clamp length:sizeof(clamp) atIndex:7];
        [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(
                ((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
                    (NSUInteger)mv_dispatch.nr0,
                (NSUInteger)n_tok,
                1)
             threadsPerThreadgroup:
                MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "shared expert scalar-row fused gate/up")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_matmul_f16_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
        const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal F16 tensor matmul received undersized activation buffers\n");
            return 0;
        }

        const uint64_t row_bytes = in_dim * sizeof(uint16_t);
        const uint64_t weight_bytes = row_bytes * out_dim;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal F16 tensor matmul range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_f32_decode_model_range(model_map,
                                                model_size,
                                                weight_offset,
                                                weight_bytes,
                                                n_tok,
                                                &inner_offset);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (n_tok == 1) {
            ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim);
            ds4_gpu_mv_dispatch mv_dispatch =
                ds4_gpu_make_plain_mv_dispatch(in_dim, 0);
            if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) {
                mv_dispatch.nr0 = 4;
                mv_dispatch.smem = 32u * 4u * sizeof(float);
            }
            mv_args.nr0 = mv_dispatch.nr0;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
            if (!pipeline) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            if (mv_dispatch.smem) {
                [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
            }
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
                                                  1,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matvec")) return 0;
            return 1;
        }

        if (n_tok <= 8 && (in_dim % 128u) == 0) {
            const int16_t nsg = 2;
            const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
            const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
            const char *fn_name = ds4_gpu_mv_ext_name(0, r1ptg);
            id<MTLComputePipelineState> pipeline =
                fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
            if (!pipeline) return 0;

            const int16_t nypsg = 32 / nxpsg;
            const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
            ds4_gpu_mul_mv_ext_args args =
                ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(uint16_t), row_bytes);

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
                                                  ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor mul_mv_ext")) return 0;
            return 1;
        }

        /*
         * Same direct-RHS TensorOps structure as Q8_0, but for F16 model
         * matrices.  The 128-token RHS tile is kept when the batch alignment
         * allows it because the later tile_n=64 retest was neutral/slower.
         */
        if (ds4_gpu_mpp_available() &&
            n_tok >= 32u &&
            (in_dim % 32u) == 0 &&
            (out_dim % 64u) == 0 &&
            (n_tok % 32u) == 0) {
            uint64_t nax_tile_n = 32u;
            if ((n_tok % 128u) == 0) {
                nax_tile_n = 128u;
            } else if ((n_tok % 64u) == 0) {
                nax_tile_n = 64u;
            }
            const char *nax_fn = nax_tile_n == 128u
                ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n128"
                : (nax_tile_n == 64u
                    ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n64"
                    : "kernel_mul_mm_f16_f32_mpp_direct_rhs");
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false);
            if (pipeline) {
                ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

                id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
                [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
                [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0];
                [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n),
                                                      (NSUInteger)out_dim / 64u,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);

                if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 NAX tensor matmul")) {
                    return 0;
                }
                return 1;
            }
            ds4_gpu_warn_mpp_fallback();
        }

        const bool bc_inp = (in_dim % 32u) != 0;
        const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_f16_f32", bc_inp, bc_out);
        if (!pipeline) return 0;

        ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u,
                                              ((NSUInteger)out_dim + 63u) / 64u,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matmul")) return 0;
    }

    return 1;
}

int ds4_gpu_matmul_f16_pair_tensor(
        ds4_gpu_tensor       *out_a,
        ds4_gpu_tensor       *out_b,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_a_offset,
        uint64_t                weight_b_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok != 1 || (in_dim & 3u) != 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outabuf = ds4_gpu_tensor_buffer(out_a);
        id<MTLBuffer> outbbuf = ds4_gpu_tensor_buffer(out_b);
        const uint64_t x_bytes = in_dim * sizeof(float);
        const uint64_t out_bytes = out_dim * sizeof(float);
        if (!xbuf || !outabuf || !outbbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out_a) < out_bytes ||
            ds4_gpu_tensor_bytes(out_b) < out_bytes) {
            fprintf(stderr, "ds4: Metal F16 paired matvec received undersized activation buffers\n");
            return 0;
        }

        const uint64_t row_bytes = in_dim * sizeof(uint16_t);
        const uint64_t weight_bytes = row_bytes * out_dim;
        if (weight_a_offset > model_size || weight_bytes > model_size - weight_a_offset ||
            weight_b_offset > model_size || weight_bytes > model_size - weight_b_offset) {
            fprintf(stderr, "ds4: Metal F16 paired matvec range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_a = 0;
        uint64_t inner_b = 0;
        id<MTLBuffer> wabuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                         weight_a_offset, weight_bytes,
                                                         &inner_a);
        id<MTLBuffer> wbbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                         weight_b_offset, weight_bytes,
                                                         &inner_b);
        if (!wabuf || !wbbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim);
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0);
        if (ds4_gpu_use_compressor_pair_nr4() &&
            (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) {
            mv_dispatch.nr0 = 4;
            mv_dispatch.smem = 32u * 4u * sizeof(float);
        }
        mv_args.nr0 = mv_dispatch.nr0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_f16_f32_pair_4", mv_dispatch.nsg);
        if (!pipeline) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
        [enc setBuffer:wabuf offset:(NSUInteger)inner_a atIndex:1];
        [enc setBuffer:wbbuf offset:(NSUInteger)inner_b atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:outabuf offset:ds4_gpu_tensor_offset(out_a) atIndex:4];
        [enc setBuffer:outbbuf offset:ds4_gpu_tensor_offset(out_b) atIndex:5];
        if (mv_dispatch.smem) {
            [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
                                              1,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 paired matvec")) return 0;
    }

    return 1;
}

int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
        ds4_gpu_tensor       *out_kv,
        ds4_gpu_tensor       *out_score,
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_kv_offset,
        uint64_t                weight_score_offset,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint64_t                in_dim,
        uint32_t                width,
        const ds4_gpu_tensor *x,
        uint32_t                ratio,
        uint32_t                pos) {
    if (!g_initialized && !ds4_gpu_init()) return -1;
    const bool force =
        getenv("DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE") != NULL;
    if ((g_quality_mode ||
         (!ds4_gpu_device_name_contains("M3") &&
          !ds4_gpu_device_name_contains("M5") && !force)) ||
        getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE") != NULL ||
        getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL ||
        getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) {
        return 0;
    }
    if (!out_kv || !out_score || !state_kv || !state_score || !model_map || !x ||
        in_dim == 0 || width == 0 || ratio == 0 ||
        (ape_type != 0u && ape_type != 1u)) {
        return -1;
    }
    if (in_dim != 4096u ||
        (width != 256u && width != 512u && width != 1024u) ||
        (ratio != 4u && ratio != 128u)) {
        return 0;
    }

    @autoreleasepool {
        const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio;
        const uint64_t row_bytes = in_dim * sizeof(uint16_t);
        const uint64_t weight_bytes = row_bytes * width;
        const uint64_t out_bytes = (uint64_t)width * sizeof(float);
        const uint64_t state_bytes =
            (uint64_t)state_rows * width * sizeof(float);
        const uint64_t ape_elem = ape_type == 1u ? sizeof(uint16_t) : sizeof(float);
        const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem;
        if (weight_kv_offset > model_size ||
            weight_bytes > model_size - weight_kv_offset ||
            weight_score_offset > model_size ||
            weight_bytes > model_size - weight_score_offset ||
            ape_offset > model_size || ape_bytes > model_size - ape_offset) {
            return -1;
        }

        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outkvbuf = ds4_gpu_tensor_buffer(out_kv);
        id<MTLBuffer> outscorebuf = ds4_gpu_tensor_buffer(out_score);
        id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
        id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
        if (!xbuf || !outkvbuf || !outscorebuf || !statekvbuf || !statescbuf ||
            ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) ||
            ds4_gpu_tensor_bytes(out_kv) < out_bytes ||
            ds4_gpu_tensor_bytes(out_score) < out_bytes ||
            ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
            ds4_gpu_tensor_bytes(state_score) < state_bytes) {
            return -1;
        }

        uint64_t weight_kv_inner = 0;
        uint64_t weight_score_inner = 0;
        uint64_t ape_inner = 0;
        id<MTLBuffer> weightkvbuf = ds4_gpu_wrap_model_range(
            model_map, model_size, weight_kv_offset, weight_bytes,
            &weight_kv_inner);
        id<MTLBuffer> weightscorebuf = ds4_gpu_wrap_model_range(
            model_map, model_size, weight_score_offset, weight_bytes,
            &weight_score_inner);
        id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(
            model_map, model_size, ape_offset, ape_bytes, &ape_inner);
        if (!weightkvbuf || !weightscorebuf || !apebuf) return -1;

        ds4_gpu_f16_matvec_args mv_args =
            ds4_gpu_make_f16_mv_args(in_dim, width);
        ds4_gpu_mv_dispatch mv_dispatch =
            ds4_gpu_make_plain_mv_dispatch(in_dim, 0);
        if (ds4_gpu_use_compressor_pair_nr4() &&
            (width == 512u || width == 1024u) && in_dim >= 4096u) {
            mv_dispatch.nr0 = 4;
            mv_dispatch.smem = 32u * 4u * sizeof(float);
        }
        mv_args.nr0 = mv_dispatch.nr0;
        ds4_gpu_dsv4_compressor_store_one_args store_args = {
            .width = width,
            .ratio = ratio,
            .pos = pos,
            .ape_type = ape_type,
        };
        id<MTLComputePipelineState> pipeline = ds4_gpu_get_mul_mv_pipeline(
            "kernel_mul_mv_f16_f32_pair_compressor_store_4",
            mv_dispatch.nsg);
        if (!pipeline) return -1;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return -1;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
        [enc setBytes:&store_args length:sizeof(store_args) atIndex:1];
        [enc setBuffer:weightkvbuf offset:(NSUInteger)weight_kv_inner atIndex:2];
        [enc setBuffer:weightscorebuf offset:(NSUInteger)weight_score_inner atIndex:3];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
        [enc setBuffer:outkvbuf offset:ds4_gpu_tensor_offset(out_kv) atIndex:5];
        [enc setBuffer:outscorebuf offset:ds4_gpu_tensor_offset(out_score) atIndex:6];
        [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:7];
        [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:8];
        [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:9];
        if (mv_dispatch.smem) {
            [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(
                ((NSUInteger)width + (NSUInteger)mv_dispatch.nr0 - 1u) /
                    (NSUInteger)mv_dispatch.nr0,
                1,
                1)
             threadsPerThreadgroup:MTLSizeMake(
                32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(
                cb, owned, "F16 paired matvec compressor state store")) {
            return -1;
        }
    }

    return 1;
}

int ds4_gpu_matmul_f32_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                n_tok) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok == 0 || n_tok > UINT32_MAX) return 0;
    if (in_dim > UINT64_MAX / n_tok / sizeof(float) ||
        out_dim > UINT64_MAX / n_tok / sizeof(float)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
        const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal F32 tensor matmul received undersized activation buffers\n");
            return 0;
        }

        const uint64_t row_bytes = in_dim * sizeof(float);
        const uint64_t weight_bytes = row_bytes * out_dim;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal F32 tensor matmul range is outside the mapped model\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf =
            ds4_gpu_wrap_f32_decode_model_range(model_map,
                                                model_size,
                                                weight_offset,
                                                weight_bytes,
                                                n_tok,
                                                &inner_offset);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (n_tok == 1) {
            ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, 1);
            ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1);
            mv_args.nr0 = mv_dispatch.nr0;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
            if (!pipeline) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            if (mv_dispatch.smem) {
                [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
            }
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
                                                  1,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matvec")) return 0;
            return 1;
        }

        if (n_tok <= 8 && (in_dim % 128u) == 0) {
            const int16_t nsg = 2;
            const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
            const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
            const char *fn_name = ds4_gpu_mv_ext_f32_name(r1ptg);
            id<MTLComputePipelineState> pipeline =
                fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
            if (!pipeline) return 0;

            const int16_t nypsg = 32 / nxpsg;
            const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
            ds4_gpu_mul_mv_ext_args args =
                ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(float), row_bytes);

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
                                                  ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor mul_mv_ext")) return 0;
            return 1;
        }

        /* Generic multi-row path (GLM prefill shapes: one grid row per
         * token through the plain matvec pipeline). */
        {
            ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, n_tok);
            ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1);
            mv_args.nr0 = mv_dispatch.nr0;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
            if (!pipeline) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
            [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
            if (mv_dispatch.smem) {
                [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
            }
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
                                                  (NSUInteger)n_tok,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matmul")) return 0;
        }
    }

    return 1;
}

int ds4_gpu_repeat_hc_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *row,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !row || n_embd == 0 || n_hc == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> rowbuf = ds4_gpu_tensor_buffer(row);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float);
        const uint64_t out_bytes = row_bytes * n_hc;
        if (!rowbuf || !outbuf ||
            ds4_gpu_tensor_bytes(row) < row_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal HC repeat received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;
        if (!ds4_gpu_encode_repeat_hc_embedding(cb,
                                                  rowbuf,
                                                  ds4_gpu_tensor_offset(row),
                                                  outbuf,
                                                  ds4_gpu_tensor_offset(out),
                                                  1,
                                                  n_embd,
                                                  n_hc)) {
            return 0;
        }
        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC repeat")) return 0;
    }

    return 1;
}

int ds4_gpu_repeat_hc_rows_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *rows,
        uint32_t                n_tokens,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> rowsbuf = ds4_gpu_tensor_buffer(rows);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float);
        const uint64_t rows_bytes = (uint64_t)n_tokens * row_bytes;
        const uint64_t out_bytes = rows_bytes * n_hc;
        if (!rowsbuf || !outbuf ||
            ds4_gpu_tensor_bytes(rows) < rows_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal HC row repeat received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;
        if (!ds4_gpu_encode_repeat_hc_embedding(cb,
                                                  rowsbuf,
                                                  ds4_gpu_tensor_offset(rows),
                                                  outbuf,
                                                  ds4_gpu_tensor_offset(out),
                                                  n_tokens,
                                                  n_embd,
                                                  n_hc)) {
            return 0;
        }
        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC row repeat")) return 0;
    }

    return 1;
}

int ds4_gpu_rms_norm_plain_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *x,
        uint32_t                n,
        float                   eps) {
    return ds4_gpu_rms_norm_plain_rows_tensor(out, x, n, 1, eps);
}

int ds4_gpu_rms_norm_plain_rows_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *x,
        uint32_t                n,
        uint32_t                rows,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (n == 0 || rows == 0 || (n & 3u) != 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t bytes = (uint64_t)n * rows * sizeof(float);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < bytes ||
            ds4_gpu_tensor_bytes(out) < bytes) {
            fprintf(stderr, "ds4: Metal plain RMS norm received undersized activation buffers\n");
            return 0;
        }

        ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps);
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_rms_norm_plain_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "plain RMS norm")) return 0;
    }

    return 1;
}

static int ds4_gpu_hc_rms_scale_project_mode(
        uint32_t in_dim,
        uint32_t out_dim,
        uint32_t n_rows) {
    const bool hard_shape = n_rows > 8u &&
        (in_dim == 16384u || in_dim == 28672u) && out_dim == 24u;
    if (!hard_shape) return 0;

    const bool force =
        getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL;
    if (getenv("DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ") != NULL ||
        (!ds4_gpu_device_name_contains("M3") && !force)) {
        return 0;
    }
    if (g_rms_norm_scale_pipeline == nil) {
        return force ? -1 : 0;
    }
    return 1;
}

int ds4_gpu_hc_rms_scale_project_f16_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *scale_scratch,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                in_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *x,
        uint32_t                n_rows,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !scale_scratch || !model_map || !x ||
        in_dim == 0u || out_dim == 0u || n_rows == 0u ||
        n_rows > INT32_MAX) {
        return 0;
    }

    const int mode =
        ds4_gpu_hc_rms_scale_project_mode(in_dim, out_dim, n_rows);
    if (mode < 0) return 0;
    if (mode == 0) {
        return ds4_gpu_rms_norm_plain_rows_tensor(
                   scale_scratch, x, in_dim, n_rows, eps) != 0 &&
               ds4_gpu_matmul_f16_tensor(
                   out, model_map, model_size, weight_offset,
                   in_dim, out_dim, scale_scratch, n_rows) != 0;
    }

    @autoreleasepool {
        const uint64_t x_elems = (uint64_t)in_dim * n_rows;
        const uint64_t out_elems = (uint64_t)out_dim * n_rows;
        if (x_elems > UINT64_MAX / sizeof(float) ||
            out_elems > UINT64_MAX / sizeof(float)) {
            return 0;
        }
        const uint64_t x_bytes = x_elems * sizeof(float);
        const uint64_t out_bytes = out_elems * sizeof(float);
        const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float);

        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> scalebuf = ds4_gpu_tensor_buffer(scale_scratch);
        if (!xbuf || !outbuf || !scalebuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(scale_scratch) < scale_bytes) {
            fprintf(stderr, "ds4: Metal HC RMS scale projection received undersized activation buffers\n");
            return 0;
        }

        const uint64_t weight_row_bytes = (uint64_t)in_dim * sizeof(uint16_t);
        if ((uint64_t)out_dim > UINT64_MAX / weight_row_bytes) return 0;
        const uint64_t weight_bytes = (uint64_t)out_dim * weight_row_bytes;
        if (weight_offset > model_size ||
            weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal HC RMS scale projection weight range is outside the mapped model\n");
            return 0;
        }

        const bool bc_inp = (in_dim % 32u) != 0u;
        const bool bc_out = (out_dim % 64u) != 0u ||
                            (n_rows % 32u) != 0u;
        id<MTLComputePipelineState> mm_pipeline =
            ds4_gpu_get_mul_mm_pipeline(
                "kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out);
        if (!mm_pipeline) {
            if (getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL) {
                return 0;
            }
            return ds4_gpu_rms_norm_plain_rows_tensor(
                       scale_scratch, x, in_dim, n_rows, eps) != 0 &&
                   ds4_gpu_matmul_f16_tensor(
                       out, model_map, model_size, weight_offset,
                       in_dim, out_dim, scale_scratch, n_rows) != 0;
        }

        id<MTLComputePipelineState> scale_pipeline = ds4_gpu_hot_pipeline(
            g_rms_norm_scale_pipeline, "kernel_rms_norm_scale_f32_4");
        if (!scale_pipeline) return 0;

        uint64_t weight_inner = 0;
        id<MTLBuffer> weightbuf = ds4_gpu_wrap_model_range(
            model_map, model_size, weight_offset, weight_bytes, &weight_inner);
        if (!weightbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_rms_norm_args norm_args =
            ds4_gpu_make_rms_norm_args(in_dim, n_rows, eps);
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:scale_pipeline];
        [enc setBytes:&norm_args length:sizeof(norm_args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setBuffer:scalebuf
                offset:ds4_gpu_tensor_offset(scale_scratch)
               atIndex:2];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(
                 ds4_gpu_rms_norm_threads(in_dim), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        ds4_gpu_mul_mm_args mm_args = ds4_gpu_make_mm_args(
            in_dim, out_dim, n_rows, weight_row_bytes);
        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:mm_pipeline];
        [enc setBytes:&mm_args length:sizeof(mm_args) atIndex:0];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setBuffer:scalebuf
                offset:ds4_gpu_tensor_offset(scale_scratch)
               atIndex:4];
        [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(
                                          ((NSUInteger)n_rows + 31u) / 32u,
                                          ((NSUInteger)out_dim + 63u) / 64u,
                                          1)
             threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(
                cb, owned, "HC RMS scale F16 projection")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_rms_norm_weight_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *x,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                n,
        float                   eps) {
    return ds4_gpu_rms_norm_weight_rows_tensor(out, x, model_map, model_size, weight_offset, n, 1, eps);
}

int ds4_gpu_rms_norm_weight_rows_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *x,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                n,
        uint32_t                rows,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (n == 0 || rows == 0 || (n & 3u) != 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t row_bytes = (uint64_t)n * sizeof(float);
        const uint64_t bytes = row_bytes * rows;
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < bytes ||
            ds4_gpu_tensor_bytes(out) < bytes) {
            fprintf(stderr, "ds4: Metal weighted RMS norm received undersized activation buffers\n");
            return 0;
        }
        if (weight_offset > model_size || row_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal weighted RMS norm range is outside the mapped model\n");
            return 0;
        }

        const bool exact_decode_weight_view =
            rows == 1u &&
            row_bytes <= (1ull << 20) &&
            getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL &&
            getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL;
        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = exact_decode_weight_view ?
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           weight_offset,
                                           row_bytes,
                                           &inner_offset) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     row_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps);
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_rms_norm_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "weighted RMS norm")) return 0;
    }

    return 1;
}

int ds4_gpu_add_rms_norm_weight_tensor(
        ds4_gpu_tensor       *norm_out,
        ds4_gpu_tensor       *sum_out,
        const ds4_gpu_tensor *a,
        const ds4_gpu_tensor *b,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                n,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!norm_out || !sum_out || !a || !b || n == 0 || (n & 3u) != 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
        id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
        id<MTLBuffer> sumbuf = ds4_gpu_tensor_buffer(sum_out);
        id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
        const uint64_t row_bytes = (uint64_t)n * sizeof(float);
        if (!abuf || !bbuf || !sumbuf || !normbuf ||
            ds4_gpu_tensor_bytes(a) < row_bytes ||
            ds4_gpu_tensor_bytes(b) < row_bytes ||
            ds4_gpu_tensor_bytes(sum_out) < row_bytes ||
            ds4_gpu_tensor_bytes(norm_out) < row_bytes) {
            fprintf(stderr, "ds4: Metal add+RMS norm received undersized activation buffers\n");
            return 0;
        }
        if (weight_offset > model_size || row_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal add+RMS norm range is outside the mapped model\n");
            return 0;
        }

        const bool exact_decode_weight_view =
            row_bytes <= (1ull << 20) &&
            getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL &&
            getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL;
        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = exact_decode_weight_view ?
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           weight_offset,
                                           row_bytes,
                                           &inner_offset) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     row_bytes,
                                     &inner_offset);
        if (!wbuf) return 0;

        ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, 1, eps);
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_add_rms_norm_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
        [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:3];
        [enc setBuffer:sumbuf offset:ds4_gpu_tensor_offset(sum_out) atIndex:4];
        [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "add+RMS norm")) return 0;
    }

    return 1;
}

int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
        ds4_gpu_tensor       *q_out,
        const ds4_gpu_tensor *q,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                q_weight_offset,
        uint32_t                q_n,
        ds4_gpu_tensor       *kv_out,
        const ds4_gpu_tensor *kv,
        uint64_t                kv_weight_offset,
        uint32_t                kv_n,
        uint32_t                rows,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!q_out || !q || !kv_out || !kv || q_n == 0 || kv_n == 0 || rows == 0 ||
        (q_n & 3u) != 0 || (kv_n & 3u) != 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> qoutbuf = ds4_gpu_tensor_buffer(q_out);
        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
        id<MTLBuffer> kvoutbuf = ds4_gpu_tensor_buffer(kv_out);

        const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float);
        const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float);
        if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf ||
            ds4_gpu_tensor_bytes(q) < q_row_bytes * rows ||
            ds4_gpu_tensor_bytes(q_out) < q_row_bytes * rows ||
            ds4_gpu_tensor_bytes(kv) < kv_row_bytes * rows ||
            ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes * rows) {
            fprintf(stderr, "ds4: Metal fused q/kv RMS norm received undersized activation buffers\n");
            return 0;
        }
        if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset ||
            kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) {
            fprintf(stderr, "ds4: Metal fused q/kv RMS norm weight range is outside the mapped model\n");
            return 0;
        }

        uint64_t q_inner_offset = 0;
        uint64_t kv_inner_offset = 0;
        id<MTLBuffer> q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                          q_weight_offset, q_row_bytes,
                                                          &q_inner_offset);
        if (!q_wbuf) return 0;
        id<MTLBuffer> kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                           kv_weight_offset, kv_row_bytes,
                                                           &kv_inner_offset);
        if (!kv_wbuf) return 0;

        ds4_gpu_qkv_rms_norm_args args = {
            .q_n = (int32_t)q_n,
            .q_n4 = (int32_t)(q_n / 4u),
            .kv_n = (int32_t)kv_n,
            .kv_n4 = (int32_t)(kv_n / 4u),
            .q_row_stride = q_row_bytes,
            .kv_row_stride = kv_row_bytes,
            .eps = eps,
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_qkv_rms_norm_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:2];
        [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3];
        [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:4];
        [enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:5];
        [enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:6];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(rows, 2, 1)
             threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv RMS norm")) return 0;
    }

    return 1;
}

int ds4_gpu_head_rms_norm_tensor(
        ds4_gpu_tensor *x,
        uint32_t          n_tok,
        uint32_t          n_head,
        uint32_t          head_dim,
        float             eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || (head_dim & 3u) != 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float);
        if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
            fprintf(stderr, "ds4: Metal head RMS norm received undersized activation buffer\n");
            return 0;
        }

        ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_3d_args(head_dim, n_head, n_tok, eps);

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_rms_norm_plain_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1)
             threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_pipeline_threads(head_dim, g_rms_norm_plain_pipeline), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "head RMS norm")) return 0;
    }

    return 1;
}

int ds4_gpu_rope_tail_tensor(
        ds4_gpu_tensor *x,
        uint32_t          n_tok,
        uint32_t          n_head,
        uint32_t          head_dim,
        uint32_t          n_rot,
        uint32_t          pos0,
        uint32_t          n_ctx_orig,
        bool              inverse,
        float             freq_base,
        float             freq_scale,
        float             ext_factor,
        float             attn_factor,
        float             beta_fast,
        float             beta_slow) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || n_rot > head_dim || (n_rot & 1u) != 0) {
        return 0;
    }
    if (n_rot == 0) return 1;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float);
        if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
            fprintf(stderr, "ds4: Metal RoPE received undersized activation buffer\n");
            return 0;
        }

        ds4_gpu_rope_tail_batch_args args = ds4_gpu_make_rope_tail_args(
            n_tok, n_head, head_dim, n_rot, n_ctx_orig, inverse,
            freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_rope_tail_inplace(cb,
                                                xbuf,
                                                ds4_gpu_tensor_offset(x),
                                                &args,
                                                n_tok,
                                                n_head,
                                                head_dim,
                                                pos0,
                                                1)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "RoPE tail")) return 0;
    }

    return 1;
}

int ds4_gpu_head_rms_norm_rope_tail_tensor(
        ds4_gpu_tensor *x,
        uint32_t          n_tok,
        uint32_t          n_head,
        uint32_t          head_dim,
        uint32_t          n_rot,
        uint32_t          pos0,
        uint32_t          n_ctx_orig,
        bool              inverse,
        float             freq_base,
        float             freq_scale,
        float             ext_factor,
        float             attn_factor,
        float             beta_fast,
        float             beta_slow,
        float             eps) {
    (void)x; (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot;
    (void)pos0; (void)n_ctx_orig; (void)inverse; (void)freq_base;
    (void)freq_scale; (void)ext_factor; (void)attn_factor;
    (void)beta_fast; (void)beta_slow; (void)eps;
    return 0;
}

int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *q_half,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint64_t              in_dim,
        uint64_t              out_dim,
        const ds4_gpu_tensor *x,
        uint32_t              n_tok,
        uint32_t              n_head,
        uint32_t              head_dim,
        uint32_t              n_rot,
        uint32_t              pos0,
        uint32_t              n_ctx_orig,
        bool                  inverse,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow,
        float                 eps) {
    (void)out; (void)q_half; (void)model_map; (void)model_size;
    (void)weight_offset; (void)in_dim; (void)out_dim; (void)x;
    (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0;
    (void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale;
    (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow;
    (void)eps;
    return 0;
}

int ds4_gpu_dsv4_fp8_kv_quantize_tensor(
        ds4_gpu_tensor *x,
        uint32_t          n_tok,
        uint32_t          head_dim,
        uint32_t          n_rot) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!x || n_tok == 0 || head_dim == 0 || n_rot > head_dim) return 0;
    if (n_rot == head_dim) return 1;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        const uint64_t bytes = (uint64_t)n_tok * head_dim * sizeof(float);
        if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
            fprintf(stderr, "ds4: Metal DSV4 FP8 KV quantize received undersized activation buffer\n");
            return 0;
        }

        ds4_gpu_dsv4_fp8_kv_quantize_args args = {
            .ne00 = head_dim,
            .ne01 = n_tok,
            .ne02 = 1,
            .ne03 = 1,
            .nb00 = sizeof(float),
            .nb01 = (uint64_t)head_dim * sizeof(float),
            .nb02 = (uint64_t)n_tok * head_dim * sizeof(float),
            .nb03 = (uint64_t)n_tok * head_dim * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)head_dim * sizeof(float),
            .nb2 = (uint64_t)n_tok * head_dim * sizeof(float),
            .nb3 = (uint64_t)n_tok * head_dim * sizeof(float),
            .n_rot = (int32_t)n_rot,
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_fp8_kv_quantize_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
        [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(n_tok, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 FP8 KV quantize")) return 0;
    }

    return 1;
}

int ds4_gpu_dsv4_indexer_qat_tensor(
        ds4_gpu_tensor *x,
        uint32_t          n_rows,
        uint32_t          head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!x || n_rows == 0 || head_dim != 128u) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        const uint64_t bytes = (uint64_t)n_rows * head_dim * sizeof(float);
        if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
            fprintf(stderr, "ds4: Metal DSV4 indexer QAT received undersized activation buffer\n");
            return 0;
        }

        ds4_gpu_dsv4_indexer_qat_args args = {
            .n_rows = n_rows,
            .head_dim = head_dim,
            .row_stride = (uint64_t)head_dim * sizeof(float),
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_indexer_qat_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 indexer Hadamard+FP4")) return 0;
    }

    return 1;
}

static void ds4_gpu_set_rows_thread_shape(
        uint32_t    width,
        NSUInteger *nth_out,
        NSUInteger *nrptg_out) {
    const NSUInteger nk0 = width ? (NSUInteger)width : 1u;
    const NSUInteger max_threads = g_set_rows_f32_i32_pipeline
        ? (NSUInteger)g_set_rows_f32_i32_pipeline.maxTotalThreadsPerThreadgroup
        : 1024u;

    NSUInteger nth = 32u;
    while (nth < nk0 && nth < max_threads) {
        nth *= 2u;
    }

    NSUInteger nrptg = 1u;
    if (nth > nk0) {
        nrptg = (nth + nk0 - 1u) / nk0;
        nth = nk0;
        if (nrptg * nth > max_threads) {
            nrptg--;
        }
    }

    if (nth > nk0) nth = nk0;
    if (nth == 0u) nth = 1u;
    if (nrptg == 0u) nrptg = 1u;

    *nth_out = nth;
    *nrptg_out = nrptg;
}

static int ds4_gpu_encode_f16_round_copy_for_raw_store(
        id<MTLCommandBuffer>   cb,
        const ds4_gpu_tensor *src,
        uint32_t               n) {
    id<MTLBuffer> srcbuf = ds4_gpu_tensor_buffer(src);
    const uint64_t src_bytes = (uint64_t)n * sizeof(float);
    if (!srcbuf || ds4_gpu_tensor_bytes(src) < src_bytes) {
        fprintf(stderr, "ds4: Metal raw KV store received undersized source buffer\n");
        return 0;
    }
    if (!ds4_gpu_ensure_scratch_buffer(&g_f16_round_scratch_buffer,
                                         &g_f16_round_scratch_bytes,
                                         (NSUInteger)n * sizeof(uint16_t),
                                         "ds4_f16_round_scratch") ||
        !ds4_gpu_ensure_scratch_buffer(&g_raw_store_round_buffer,
                                         &g_raw_store_round_bytes,
                                         (NSUInteger)n * sizeof(float),
                                         "ds4_raw_store_round")) {
        return 0;
    }

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         srcbuf,
                                         ds4_gpu_tensor_offset(src),
                                         g_f16_round_scratch_buffer,
                                         0,
                                         n)) {
        return 0;
    }
    return ds4_gpu_encode_cpy_f16_f32_1d(cb,
                                           g_f16_round_scratch_buffer,
                                           0,
                                           g_raw_store_round_buffer,
                                           0,
                                           n);
}

static int ds4_gpu_encode_set_rows_f32_i32(
        id<MTLCommandBuffer> cb,
        ds4_gpu_tensor    *dst,
        id<MTLBuffer>        srcbuf,
        NSUInteger           src_off,
        const int32_t       *rows,
        uint32_t             n_rows,
        uint32_t             dst_rows,
        uint32_t             width) {
    id<MTLBuffer> dstbuf = ds4_gpu_tensor_buffer(dst);
    const uint64_t dst_bytes = (uint64_t)dst_rows * width * sizeof(float);
    const uint64_t src_bytes = (uint64_t)n_rows * width * sizeof(float);
    if (!dstbuf || !srcbuf || !rows || n_rows == 0 || width == 0 ||
        ds4_gpu_tensor_bytes(dst) < dst_bytes ||
        src_bytes > NSUIntegerMax - src_off) {
        fprintf(stderr, "ds4: Metal DS4 set_rows received invalid buffers\n");
        return 0;
    }

    const uint64_t row_bytes = (uint64_t)width * sizeof(float);
    const uint64_t rows_bytes = (uint64_t)n_rows * sizeof(int32_t);
    ds4_gpu_set_rows_args args = {
        .nk0 = (int32_t)width,
        .ne01 = (int32_t)n_rows,
        .nb01 = row_bytes,
        .nb02 = (uint64_t)n_rows * row_bytes,
        .nb03 = (uint64_t)n_rows * row_bytes,
        .ne11 = 1,
        .ne12 = 1,
        .nb10 = sizeof(int32_t),
        .nb11 = rows_bytes,
        .nb12 = rows_bytes,
        .nb1 = row_bytes,
        .nb2 = (uint64_t)dst_rows * row_bytes,
        .nb3 = (uint64_t)dst_rows * row_bytes,
    };

    NSUInteger nth;
    NSUInteger nrptg;
    ds4_gpu_set_rows_thread_shape(width, &nth, &nrptg);

    id<MTLBuffer> rowsbuf = nil;
    if (rows_bytes > 4096u) {
        rowsbuf = ds4_gpu_new_transient_buffer((NSUInteger)rows_bytes, "ds4_set_rows_indices");
        if (!rowsbuf) return 0;
        memcpy([rowsbuf contents], rows, (NSUInteger)rows_bytes);
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_set_rows_f32_i32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:srcbuf offset:src_off atIndex:1];
    if (rowsbuf) {
        [enc setBuffer:rowsbuf offset:0 atIndex:2];
    } else {
        [enc setBytes:rows length:(NSUInteger)rows_bytes atIndex:2];
    }
    [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + nrptg - 1u) / nrptg, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, nrptg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_add_f32_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        a,
        NSUInteger           a_off,
        id<MTLBuffer>        b,
        NSUInteger           b_off,
        id<MTLBuffer>        out,
        NSUInteger           out_off,
        uint32_t             n) {
    if (!cb || !a || !b || !out || n == 0) return 0;

    ds4_gpu_add_flat_args args = { .n = n };
    NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > 256u) nth = 256u;
    if (nth > (NSUInteger)n) nth = (NSUInteger)n;
    if (nth == 0u) nth = 1u;
    const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_add2_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:a offset:a_off atIndex:1];
    [enc setBuffer:b offset:b_off atIndex:2];
    [enc setBuffer:out offset:out_off atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

int ds4_gpu_store_raw_kv_tensor(
        ds4_gpu_tensor       *raw_cache,
        const ds4_gpu_tensor *kv,
        uint32_t                raw_cap,
        uint32_t                row,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!raw_cache || !kv || raw_cap == 0 || row >= raw_cap || head_dim == 0 || raw_cap > INT32_MAX) return 0;

    @autoreleasepool {
        const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
        if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
            fprintf(stderr, "ds4: Metal raw KV store received undersized destination buffer\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const int32_t row_i32 = (int32_t)row;
        if (!ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, head_dim) ||
            !ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache,
                                               g_raw_store_round_buffer,
                                               0,
                                               &row_i32,
                                               1,
                                               raw_cap,
                                               head_dim)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV DS4 set_rows store")) return 0;
    }

    return 1;
}

/* Release decode fused KV finalizer.  Reference paths are selected by the C
 * graph driver; this Objective-C entry point always means "use the fused
 * Metal kernel." */
int ds4_gpu_kv_fp8_store_raw_tensor(
        ds4_gpu_tensor *kv,
        ds4_gpu_tensor *raw_cache,
        uint32_t          raw_cap,
        uint32_t          row,
        uint32_t          head_dim,
        uint32_t          n_rot) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 ||
        n_rot > head_dim || raw_cap > INT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_cache);
        const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float);
        const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
        if (!kvbuf || !rawbuf ||
            ds4_gpu_tensor_bytes(kv) < kv_bytes ||
            ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
            fprintf(stderr, "ds4: Metal fused KV FP8/raw-store received undersized buffers\n");
            return 0;
        }

        ds4_gpu_dsv4_kv_fp8_store_args args = {
            .head_dim = (int32_t)head_dim,
            .n_rot = (int32_t)n_rot,
            .raw_row = (int32_t)row,
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_kv_fp8_store_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1];
        [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:2];
        [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "KV FP8/raw-store fused")) return 0;
    }

    return 1;
}

int ds4_gpu_store_raw_kv_batch_tensor(
        ds4_gpu_tensor       *raw_cache,
        const ds4_gpu_tensor *kv,
        uint32_t                raw_cap,
        uint32_t                pos0,
        uint32_t                n_tokens,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!raw_cache || !kv || raw_cap == 0 || n_tokens == 0 || head_dim == 0 || raw_cap > INT32_MAX) return 0;

    @autoreleasepool {
        const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
        if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
            fprintf(stderr, "ds4: Metal raw KV batch store received undersized destination buffer\n");
            return 0;
        }

        int32_t rows_stack[512];
        int32_t *rows = rows_stack;
        if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) {
            rows = malloc((size_t)n_tokens * sizeof(*rows));
            if (!rows) {
                fprintf(stderr, "ds4: failed to allocate raw KV set_rows index list\n");
                return 0;
            }
        }
        for (uint32_t t = 0; t < n_tokens; t++) {
            rows[t] = (int32_t)((pos0 + t) % raw_cap);
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) {
            if (rows != rows_stack) free(rows);
            return 0;
        }

        const uint64_t n = (uint64_t)n_tokens * head_dim;
        const int ok = n <= UINT32_MAX &&
            ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, (uint32_t)n) &&
            ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache,
                                               g_raw_store_round_buffer,
                                               0,
                                               rows,
                                               n_tokens,
                                               raw_cap,
                                               head_dim);
        if (rows != rows_stack) free(rows);
        if (!ok) return 0;

        if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV batch DS4 set_rows store")) return 0;
    }

    return 1;
}

static int ds4_gpu_encode_compressor_score_with_ape(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        score_src,
        NSUInteger           score_src_offset,
        id<MTLBuffer>        score_dst,
        NSUInteger           score_dst_offset,
        id<MTLBuffer>        apebuf,
        NSUInteger           ape_offset,
        uint32_t             ape_type,
        uint32_t             width,
        uint32_t             ratio,
        uint32_t             pos0,
        uint32_t             n_tokens) {
    if (!cb || !score_src || !score_dst || !apebuf ||
        width == 0 || ratio == 0 || n_tokens == 0 ||
        (ape_type != 0u && ape_type != 1u)) {
        return 0;
    }

    const uint64_t total_elems64 = (uint64_t)n_tokens * width;
    if (total_elems64 > UINT32_MAX) {
        fprintf(stderr, "ds4: Metal compressor APE add received too many elements\n");
        return 0;
    }
    const uint32_t total_elems = (uint32_t)total_elems64;

    const bool force_fused =
        getenv("DS4_METAL_ENABLE_COMPRESSOR_APE_ADD") != NULL;
    const bool use_fused =
        (ds4_gpu_device_name_contains("M3") || force_fused) &&
        getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD") == NULL;
    if (use_fused) {
        id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
            ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16"
                           : "kernel_dsv4_compressor_score_ape_f32");
        if (pipeline) {
            ds4_gpu_dsv4_compressor_score_ape_args args = {
                .width = width,
                .ratio = ratio,
                .pos0 = pos0,
                .n_tokens = n_tokens,
            };
            NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
            if (nth > 256u) nth = 256u;
            if (nth == 0) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:score_src offset:score_src_offset atIndex:1];
            [enc setBuffer:apebuf offset:ape_offset atIndex:2];
            [enc setBuffer:score_dst offset:score_dst_offset atIndex:3];
            [enc dispatchThreadgroups:MTLSizeMake(
                    ((NSUInteger)total_elems + nth - 1u) / nth, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);
            return 1;
        }
        if (force_fused) return 0;
    }

    const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float);
    if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer,
                                         &g_compressor_store_ape_bytes,
                                         scratch_bytes,
                                         "ds4_compressor_store_ape")) {
        return 0;
    }

    const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
    uint32_t copied_rows = 0;
    uint32_t pos_mod = pos0 % ratio;
    while (copied_rows < n_tokens) {
        uint32_t seg_rows = ratio - pos_mod;
        if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows;
        const uint32_t seg_elems = seg_rows * width;
        const NSUInteger src_off = ape_offset + (NSUInteger)pos_mod * width * elem_ape;
        const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float);
        int ok;
        if (ape_type == 1u) {
            ok = ds4_gpu_encode_cpy_f16_f32_1d(cb,
                                                 apebuf,
                                                 src_off,
                                                 g_compressor_store_ape_buffer,
                                                 dst_off,
                                                 seg_elems);
        } else {
            ok = ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                                 apebuf,
                                                 src_off,
                                                 g_compressor_store_ape_buffer,
                                                 dst_off,
                                                 seg_elems);
        }
        if (!ok) return 0;
        copied_rows += seg_rows;
        pos_mod = 0;
    }

    return ds4_gpu_encode_add_f32_1d(cb,
                                       score_src,
                                       score_src_offset,
                                       g_compressor_store_ape_buffer,
                                       0,
                                       score_dst,
                                       score_dst_offset,
                                       total_elems);
}

static int ds4_gpu_encode_compressor_set_rows_projected(
        id<MTLCommandBuffer> cb,
        ds4_gpu_tensor    *state_kv,
        ds4_gpu_tensor    *state_score,
        id<MTLBuffer>        kvbuf,
        NSUInteger           kv_offset,
        id<MTLBuffer>        scorebuf,
        NSUInteger           score_offset,
        id<MTLBuffer>        apebuf,
        NSUInteger           ape_offset,
        uint32_t             ape_type,
        uint32_t             width,
        uint32_t             ratio,
        uint32_t             pos0,
        const int32_t       *rows,
        uint32_t             n_rows,
        uint32_t             state_rows) {
    if (!cb || !state_kv || !state_score || !kvbuf || !scorebuf ||
        !apebuf || !rows || width == 0 || n_rows == 0 || state_rows == 0) {
        return 0;
    }

    const NSUInteger score_scratch_bytes = (NSUInteger)n_rows * width * sizeof(float);
    if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
                                         &g_compressor_store_score_bytes,
                                         score_scratch_bytes,
                                         "ds4_compressor_store_score")) {
        return 0;
    }

    return ds4_gpu_encode_compressor_score_with_ape(cb,
                                                      scorebuf,
                                                      score_offset,
                                                      g_compressor_store_score_buffer,
                                                      0,
                                                      apebuf,
                                                      ape_offset,
                                                      ape_type,
                                                      width,
                                                      ratio,
                                                      pos0,
                                                      n_rows) &&
           ds4_gpu_encode_set_rows_f32_i32(cb,
                                             state_kv,
                                             kvbuf,
                                             kv_offset,
                                             rows,
                                             n_rows,
                                             state_rows,
                                             width) &&
           ds4_gpu_encode_set_rows_f32_i32(cb,
                                             state_score,
                                             g_compressor_store_score_buffer,
                                             0,
                                             rows,
                                             n_rows,
                                             state_rows,
                                             width);
}

static int ds4_gpu_compressor_store_one_tensor(
        const ds4_gpu_tensor *kv,
        const ds4_gpu_tensor *sc,
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint32_t                width,
        uint32_t                ratio,
        uint32_t                pos) {
    if (!kv || !sc || !state_kv || !state_score || !model_map ||
        width == 0 || ratio == 0 || (ape_type != 0u && ape_type != 1u)) {
        return 0;
    }

    id<MTLComputePipelineState> pipeline =
        ds4_gpu_hot_pipeline(g_dsv4_compressor_store_one_pipeline,
                                "kernel_dsv4_compressor_store_one");
    if (!pipeline) return 0;

    const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio;
    const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
    const uint64_t row_bytes = (uint64_t)width * sizeof(float);
    const uint64_t state_bytes = (uint64_t)state_rows * row_bytes;
    const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
    if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
        ds4_gpu_tensor_bytes(kv) < row_bytes ||
        ds4_gpu_tensor_bytes(sc) < row_bytes ||
        ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
        ds4_gpu_tensor_bytes(state_score) < state_bytes) {
        return 0;
    }

    uint64_t ape_inner = 0;
    id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                       ape_offset, ape_bytes,
                                                       &ape_inner);
    id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
    id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
    id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
    id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
    if (!apebuf || !kvbuf || !scbuf || !statekvbuf || !statescbuf) return 0;

    ds4_gpu_dsv4_compressor_store_one_args args = {
        .width = width,
        .ratio = ratio,
        .pos = pos,
        .ape_type = ape_type,
    };

    int owned = 0;
    id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
    if (!cb) return 0;

    const NSUInteger nth = 256u;
    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1];
    [enc setBuffer:scbuf offset:ds4_gpu_tensor_offset(sc) atIndex:2];
    [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:3];
    [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:4];
    [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:5];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)width + nth - 1u) / nth, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return ds4_gpu_finish_command_buffer(cb, owned, "compressor one-row store");
}

int ds4_gpu_compressor_store_batch_tensor(
        const ds4_gpu_tensor *kv,
        const ds4_gpu_tensor *sc,
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint32_t                head_dim,
        uint32_t                ratio,
        uint32_t                pos0,
        uint32_t                n_tokens) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!kv || !sc || !state_kv || !state_score || !model_map ||
        head_dim == 0 || ratio == 0 || n_tokens == 0 ||
        (ape_type != 0u && ape_type != 1u)) {
        return 0;
    }

    @autoreleasepool {
        const uint32_t coff = ratio == 4u ? 2u : 1u;
        const uint32_t width = coff * head_dim;
        const uint32_t state_rows = coff * ratio;
        const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
        const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
        const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
        const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;

        if (ape_offset > model_size || ape_bytes > model_size - ape_offset) {
            fprintf(stderr, "ds4: Metal compressor batch APE range is outside the mapped model\n");
            return 0;
        }

        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
        id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
        if (!kvbuf || !scbuf ||
            ds4_gpu_tensor_bytes(kv) < kv_bytes ||
            ds4_gpu_tensor_bytes(sc) < kv_bytes ||
            ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
            ds4_gpu_tensor_bytes(state_score) < state_bytes) {
            fprintf(stderr, "ds4: Metal compressor batch store received undersized buffers\n");
            return 0;
        }

        uint64_t ape_inner = 0;
        id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
        if (!apebuf) return 0;

        const uint64_t total_elems64 = (uint64_t)n_tokens * width;
        if (total_elems64 > UINT32_MAX || state_rows > INT32_MAX) {
            fprintf(stderr, "ds4: Metal compressor batch store received too many elements\n");
            return 0;
        }
        const uint32_t total_elems = (uint32_t)total_elems64;
        const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float);
        if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer,
                                             &g_compressor_store_ape_bytes,
                                             scratch_bytes,
                                             "ds4_compressor_store_ape") ||
            !ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
                                             &g_compressor_store_score_bytes,
                                             scratch_bytes,
                                             "ds4_compressor_store_score")) {
            return 0;
        }

        int32_t rows_stack[16];
        int32_t *rows = rows_stack;
        if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) {
            rows = malloc((size_t)n_tokens * sizeof(*rows));
            if (!rows) {
                fprintf(stderr, "ds4: failed to allocate compressor set_rows index list\n");
                return 0;
            }
        }
        for (uint32_t t = 0; t < n_tokens; t++) {
            const uint32_t pos_mod = (pos0 + t) % ratio;
            rows[t] = (int32_t)(ratio == 4u ? ratio + pos_mod : pos_mod);
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) {
            if (rows != rows_stack) free(rows);
            return 0;
        }

        int ok = 1;
        uint32_t copied_rows = 0;
        uint32_t pos_mod = pos0 % ratio;
        while (ok && copied_rows < n_tokens) {
            uint32_t seg_rows = ratio - pos_mod;
            if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows;
            const uint32_t seg_elems = seg_rows * width;
            const NSUInteger src_off = (NSUInteger)ape_inner +
                                       (NSUInteger)pos_mod * width * elem_ape;
            const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float);
            if (ape_type == 1u) {
                ok = ds4_gpu_encode_cpy_f16_f32_1d(cb,
                                                     apebuf,
                                                     src_off,
                                                     g_compressor_store_ape_buffer,
                                                     dst_off,
                                                     seg_elems);
            } else {
                ok = ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                                     apebuf,
                                                     src_off,
                                                     g_compressor_store_ape_buffer,
                                                     dst_off,
                                                     seg_elems);
            }
            copied_rows += seg_rows;
            pos_mod = 0;
        }

        if (ok) {
            ok = ds4_gpu_encode_add_f32_1d(cb,
                                             scbuf,
                                             ds4_gpu_tensor_offset(sc),
                                             g_compressor_store_ape_buffer,
                                             0,
                                             g_compressor_store_score_buffer,
                                             0,
                                             total_elems);
        }
        if (ok) {
            ok = ds4_gpu_encode_set_rows_f32_i32(cb,
                                                   state_kv,
                                                   kvbuf,
                                                   ds4_gpu_tensor_offset(kv),
                                                   rows,
                                                   n_tokens,
                                                   state_rows,
                                                   width);
        }
        if (ok) {
            ok = ds4_gpu_encode_set_rows_f32_i32(cb,
                                                   state_score,
                                                   g_compressor_store_score_buffer,
                                                   0,
                                                   rows,
                                                   n_tokens,
                                                   state_rows,
                                                   width);
        }
        if (rows != rows_stack) free(rows);
        if (!ok) return 0;

        if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor batch DS4 store")) return 0;
    }

    return 1;
}

static ds4_gpu_bin_args ds4_gpu_make_bin_contiguous_3d_args(
        uint32_t cols,
        uint32_t rows,
        uint32_t planes) {
    const uint64_t row_bytes = (uint64_t)cols * sizeof(float);
    const uint64_t plane_bytes = (uint64_t)rows * row_bytes;
    return (ds4_gpu_bin_args) {
        .ne00 = (int32_t)cols,
        .ne01 = (int32_t)rows,
        .ne02 = (int32_t)planes,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = row_bytes,
        .nb02 = plane_bytes,
        .nb03 = (uint64_t)planes * plane_bytes,
        .ne10 = (int32_t)cols,
        .ne11 = (int32_t)rows,
        .ne12 = (int32_t)planes,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = row_bytes,
        .nb12 = plane_bytes,
        .nb13 = (uint64_t)planes * plane_bytes,
        .ne0 = (int32_t)cols,
        .ne1 = (int32_t)rows,
        .ne2 = (int32_t)planes,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = row_bytes,
        .nb2 = plane_bytes,
        .nb3 = (uint64_t)planes * plane_bytes,
        .offs = 0,
        .o1 = { 0 },
    };
}

static int ds4_gpu_encode_softmax_f32_contiguous(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             width,
        uint32_t             rows,
        uint32_t             planes) {
    if (!cb || !src || !dst || width == 0 || rows == 0 || planes == 0) return 0;

    const uint64_t row_bytes = (uint64_t)width * sizeof(float);
    const uint64_t plane_bytes = (uint64_t)rows * row_bytes;
    ds4_gpu_softmax_args args = {
        .ne00 = (int32_t)width,
        .ne01 = (int32_t)rows,
        .ne02 = (int32_t)planes,
        .nb01 = row_bytes,
        .nb02 = plane_bytes,
        .nb03 = (uint64_t)planes * plane_bytes,
        .ne11 = (int32_t)width,
        .ne12 = (int32_t)rows,
        .ne13 = (int32_t)planes,
        .nb11 = row_bytes,
        .nb12 = plane_bytes,
        .nb13 = (uint64_t)planes * plane_bytes,
        .nb1 = row_bytes,
        .nb2 = plane_bytes,
        .nb3 = (uint64_t)planes * plane_bytes,
        .scale = 1.0f,
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 1,
    };

    id<MTLComputePipelineState> pipeline =
        (width % 4u) == 0 ? g_soft_max_f32_4_pipeline : g_soft_max_f32_pipeline;
    if (!pipeline) return 0;

    NSUInteger nth = 32u;
    if ((width % 4u) == 0) {
        while (nth < (NSUInteger)(width / 4u) &&
               nth * (NSUInteger)rows * (NSUInteger)planes < 256u) {
            nth *= 2u;
        }
    } else {
        while (nth < (NSUInteger)width &&
               nth * (NSUInteger)rows * (NSUInteger)planes < 256u) {
            nth *= 2u;
        }
    }
    const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:src offset:src_off atIndex:2];
    [enc setBuffer:src offset:src_off atIndex:3];
    [enc setBuffer:dst offset:dst_off atIndex:4];
    [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(rows, planes, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(
        id<MTLCommandBuffer> cb,
        ds4_gpu_tensor    *out,
        id<MTLBuffer>        kvbuf,
        NSUInteger           kv_offset,
        uint64_t             kv_nb0,
        uint64_t             kv_nb1,
        uint64_t             kv_nb2,
        id<MTLBuffer>        scorebuf,
        NSUInteger           score_offset,
        uint64_t             score_nb0,
        uint64_t             score_nb1,
        uint64_t             score_nb2,
        uint32_t             n_rows,
        uint32_t             head_dim) {
    id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
    if (!cb || !outbuf || !kvbuf || !scorebuf || n_rows == 0 || head_dim == 0 ||
        ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float)) {
        return 0;
    }

    const NSUInteger pack_bytes = (NSUInteger)n_rows * head_dim * sizeof(float);
    if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer,
                                         &g_compressor_pool_product_bytes,
                                         pack_bytes,
                                         "ds4_compressor_pool_product") ||
        !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer,
                                         &g_compressor_pool_score_cont_bytes,
                                         pack_bytes,
                                         "ds4_compressor_pool_score_cont") ||
        !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer,
                                         &g_compressor_pool_softmax_bytes,
                                         pack_bytes,
                                         "ds4_compressor_pool_softmax")) {
        return 0;
    }

    const uint64_t cont_row_stride = (uint64_t)n_rows * sizeof(float);
    const uint64_t cont_plane_stride = (uint64_t)head_dim * cont_row_stride;

    /*
     * Keep the n_comp == 1 compressor path as the unfused graph sequence:
     *
     *   score = soft_max(contiguous(score))
     *   pooled = sum_rows(contiguous(kv) * score)
     *
     * The fused DS4 pool kernel is mathematically equivalent, but it reduces in
     * a different order. That is enough to create ~1e-6 compressor differences
     * and later FP8/routing flips, so this path intentionally keeps the same
     * operation boundary and memory layout as the graph.
     */
    ds4_gpu_bin_args mul_args =
        ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1);

    return
        ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb,
                                                    kvbuf,
                                                    kv_offset,
                                                    g_compressor_pool_product_buffer,
                                                    0,
                                                    n_rows,
                                                    head_dim,
                                                    1,
                                                    kv_nb0,
                                                    kv_nb1,
                                                    kv_nb2,
                                                    cont_row_stride,
                                                    cont_plane_stride) &&
        ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb,
                                                    scorebuf,
                                                    score_offset,
                                                    g_compressor_pool_score_cont_buffer,
                                                    0,
                                                    n_rows,
                                                    head_dim,
                                                    1,
                                                    score_nb0,
                                                    score_nb1,
                                                    score_nb2,
                                                    cont_row_stride,
                                                    cont_plane_stride) &&
        ds4_gpu_encode_softmax_f32_contiguous(cb,
                                                g_compressor_pool_score_cont_buffer,
                                                0,
                                                g_compressor_pool_softmax_buffer,
                                                0,
                                                n_rows,
                                                head_dim,
                                                1) &&
        ds4_gpu_encode_bin_f32_rows(cb,
                                      g_mul_pipeline,
                                      &mul_args,
                                      g_compressor_pool_product_buffer,
                                      0,
                                      g_compressor_pool_softmax_buffer,
                                      0,
                                      g_compressor_pool_product_buffer,
                                      0) &&
        ds4_gpu_encode_sum_rows_f32(cb,
                                      g_compressor_pool_product_buffer,
                                      0,
                                      outbuf,
                                      ds4_gpu_tensor_offset(out),
                                      n_rows,
                                      head_dim);
}

static int ds4_gpu_encode_dsv4_softmax_pool(
        id<MTLCommandBuffer> cb,
        ds4_gpu_tensor    *out,
        id<MTLBuffer>        kvbuf,
        NSUInteger           kv_offset,
        uint64_t             kv_nb0,
        uint64_t             kv_nb1,
        uint64_t             kv_nb2,
        id<MTLBuffer>        scorebuf,
        NSUInteger           score_offset,
        uint64_t             score_nb0,
        uint64_t             score_nb1,
        uint64_t             score_nb2,
        uint32_t             n_rows,
        uint32_t             head_dim,
        uint32_t             n_comp) {
    id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
    if (!cb || !outbuf || !kvbuf || !scorebuf ||
        n_rows == 0 || head_dim == 0 || n_comp == 0 ||
        ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * n_comp * sizeof(float)) {
        return 0;
    }

    if (n_comp == 1) {
        return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(cb,
                                                                out,
                                                                kvbuf,
                                                                kv_offset,
                                                                kv_nb0,
                                                                kv_nb1,
                                                                kv_nb2,
                                                                scorebuf,
                                                                score_offset,
                                                                score_nb0,
                                                                score_nb1,
                                                                score_nb2,
                                                                n_rows,
                                                                head_dim);
    }

    ds4_gpu_dsv4_softmax_pool_args args = {
        .ne00 = (int64_t)n_rows,
        .ne01 = (int64_t)head_dim,
        .ne02 = (int64_t)n_comp,
        .nb00 = kv_nb0,
        .nb01 = kv_nb1,
        .nb02 = kv_nb2,
        .nb10 = score_nb0,
        .nb11 = score_nb1,
        .nb12 = score_nb2,
        .ne0 = (int64_t)head_dim,
        .ne1 = (int64_t)n_comp,
        .nb0 = sizeof(float),
        .nb1 = (uint64_t)head_dim * sizeof(float),
    };
    const uint64_t n = (uint64_t)head_dim * n_comp;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_dsv4_softmax_pool_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:kvbuf offset:kv_offset atIndex:1];
    [enc setBuffer:scorebuf offset:score_offset atIndex:2];
    [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_concat_f32_dim1(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src0,
        NSUInteger           src0_offset,
        uint32_t             src0_rows,
        uint64_t             src0_row_stride,
        id<MTLBuffer>        src1,
        NSUInteger           src1_offset,
        uint32_t             src1_rows,
        uint64_t             src1_row_stride,
        id<MTLBuffer>        dst,
        NSUInteger           dst_offset,
        uint32_t             cols,
        uint64_t             dst_row_stride) {
    if (!cb || !src0 || !src1 || !dst || cols == 0 || src0_rows == 0 || src1_rows == 0) {
        return 0;
    }

    const uint32_t rows = src0_rows + src1_rows;
    const uint64_t src0_plane = (uint64_t)src0_rows * src0_row_stride;
    const uint64_t src1_plane = (uint64_t)src1_rows * src1_row_stride;
    const uint64_t dst_plane = (uint64_t)rows * dst_row_stride;
    ds4_gpu_concat_args args = {
        .ne00 = (int32_t)cols,
        .ne01 = (int32_t)src0_rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = src0_row_stride,
        .nb02 = src0_plane,
        .nb03 = src0_plane,
        .ne10 = (int32_t)cols,
        .ne11 = (int32_t)src1_rows,
        .ne12 = 1,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = src1_row_stride,
        .nb12 = src1_plane,
        .nb13 = src1_plane,
        .ne0 = (int32_t)cols,
        .ne1 = (int32_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = dst_row_stride,
        .nb2 = dst_plane,
        .nb3 = dst_plane,
        .dim = 1,
    };

    NSUInteger nth = cols < 1024u ? (NSUInteger)cols : 1024u;
    const NSUInteger max_threads = g_concat_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_concat_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src0 offset:src0_offset atIndex:1];
    [enc setBuffer:src1 offset:src1_offset atIndex:2];
    [enc setBuffer:dst offset:dst_offset atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) {
    const bool force =
        getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL;
    const bool default_shape = head_dim == 128u || head_dim == 512u;
    if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION") != NULL ||
        (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) {
        return 0;
    }
    if (g_dsv4_compressor_pack_ratio4_pipeline == nil) {
        return force ? -1 : 0;
    }
    return 1;
}

static bool ds4_gpu_buffer_ranges_overlap(
        id<MTLBuffer> a,
        NSUInteger    a_offset,
        uint64_t      a_bytes,
        id<MTLBuffer> b,
        NSUInteger    b_offset,
        uint64_t      b_bytes) {
    if (a != b || a_bytes == 0u || b_bytes == 0u) return false;
    if (a_offset <= b_offset) {
        return a_bytes > (uint64_t)(b_offset - a_offset);
    }
    return b_bytes > (uint64_t)(a_offset - b_offset);
}

static int ds4_gpu_compressor_ratio4_direct_pool_mode(
        uint32_t head_dim,
        uint32_t n_comp) {
    // One compressed row intentionally uses the legacy GGML reduction graph.
    if (n_comp <= 1u) return 0;

    const bool force =
        getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL;
    const bool default_shape = head_dim == 128u || head_dim == 512u;
    if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL ||
        (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) {
        return 0;
    }
    if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) {
        return force ? -1 : 0;
    }
    return 1;
}

static int ds4_gpu_encode_compressor_ratio4_direct_pool(
        id<MTLCommandBuffer> cb,
        ds4_gpu_tensor    *out,
        id<MTLBuffer>        kvbuf,
        NSUInteger           kv_offset,
        id<MTLBuffer>        scorebuf,
        NSUInteger           score_offset,
        id<MTLBuffer>        statekvbuf,
        NSUInteger           state_kv_offset,
        id<MTLBuffer>        statescbuf,
        NSUInteger           state_score_offset,
        uint32_t             head_dim,
        uint32_t             n_comp,
        bool                 replay) {
    id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
    const uint64_t n = (uint64_t)head_dim * n_comp;
    if (!cb || !outbuf || !kvbuf || !scorebuf || !statekvbuf || !statescbuf ||
        head_dim == 0u || n_comp <= 1u || n > UINT32_MAX ||
        ds4_gpu_tensor_bytes(out) < n * sizeof(float)) {
        return 0;
    }

    id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
        g_dsv4_softmax_pool_ratio4_direct_pipeline,
        "kernel_dsv4_softmax_pool_ratio4_direct");
    if (!pipeline) return 0;

    ds4_gpu_dsv4_softmax_pool_ratio4_direct_args args = {
        .n_rows = 8,
        .head_dim = head_dim,
        .n_comp = n_comp,
        .replay = replay ? 1u : 0u,
        .pad = 0u,
    };

    NSUInteger nth = 256u;
    if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
        nth = pipeline.maxTotalThreadsPerThreadgroup;
    }
    if (nth == 0u) return 0;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:kvbuf offset:kv_offset atIndex:1];
    [enc setBuffer:scorebuf offset:score_offset atIndex:2];
    [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3];
    [enc setBuffer:statescbuf offset:state_score_offset atIndex:4];
    [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:5];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + nth - 1u) / nth, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_compressor_pack_ratio4(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        kvbuf,
        NSUInteger           kv_offset,
        id<MTLBuffer>        scorebuf,
        NSUInteger           score_offset,
        id<MTLBuffer>        statekvbuf,
        NSUInteger           state_kv_offset,
        id<MTLBuffer>        statescbuf,
        NSUInteger           state_score_offset,
        uint32_t             head_dim,
        uint32_t             n_comp,
        bool                 replay) {
    if (!cb || !kvbuf || !scorebuf || !statekvbuf || !statescbuf ||
        !g_compressor_pool_kv_buffer || !g_compressor_pool_score_buffer ||
        head_dim == 0 || n_comp == 0 || head_dim > UINT32_MAX / 2u) {
        return 0;
    }

    const uint64_t total_elems64 = (uint64_t)n_comp * 8u * head_dim;
    if (total_elems64 > UINT32_MAX || n_comp > UINT32_MAX / 8u) return 0;
    id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
        g_dsv4_compressor_pack_ratio4_pipeline,
        "kernel_dsv4_compressor_pack_ratio4");
    if (!pipeline) return 0;

    NSUInteger nth = head_dim;
    if (nth > 256u) nth = 256u;
    if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
        nth = pipeline.maxTotalThreadsPerThreadgroup;
    }
    if (nth == 0) return 0;

    ds4_gpu_dsv4_compressor_pack_ratio4_args args = {
        .head_dim = head_dim,
        .n_comp = n_comp,
        .replay = replay ? 1u : 0u,
        .n_threads = (uint32_t)nth,
    };

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:kvbuf offset:kv_offset atIndex:1];
    [enc setBuffer:scorebuf offset:score_offset atIndex:2];
    [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3];
    [enc setBuffer:statescbuf offset:state_score_offset atIndex:4];
    [enc setBuffer:g_compressor_pool_kv_buffer offset:0 atIndex:5];
    [enc setBuffer:g_compressor_pool_score_buffer offset:0 atIndex:6];
    [enc dispatchThreadgroups:MTLSizeMake(n_comp, 8u, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_compressor_pool(
        id<MTLCommandBuffer>   cb,
        ds4_gpu_tensor      *out,
        const ds4_gpu_tensor *state_kv,
        const ds4_gpu_tensor *state_score,
        uint32_t               head_dim,
        uint32_t               ratio) {
    id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
    id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
    if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0 || ratio == 0) return 0;

    const uint32_t coff = ratio == 4u ? 2u : 1u;
    const uint32_t width = coff * head_dim;
    const uint32_t rows = coff * ratio;
    const uint64_t state_bytes = (uint64_t)width * rows * sizeof(float);
    if (ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
        ds4_gpu_tensor_bytes(state_score) < state_bytes) {
        return 0;
    }

    if (ratio != 4u) {
        const uint64_t row_stride = (uint64_t)width * sizeof(float);
        return ds4_gpu_encode_dsv4_softmax_pool(cb,
                                                  out,
                                                  statekvbuf,
                                                  ds4_gpu_tensor_offset(state_kv),
                                                  row_stride,
                                                  sizeof(float),
                                                  (uint64_t)rows * row_stride,
                                                  statescbuf,
                                                  ds4_gpu_tensor_offset(state_score),
                                                  row_stride,
                                                  sizeof(float),
                                                  (uint64_t)rows * row_stride,
                                                  ratio,
                                                  head_dim,
                                                  1);
    }

    const NSUInteger packed_bytes = (NSUInteger)8u * head_dim * sizeof(float);
    if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer,
                                         &g_compressor_pool_kv_bytes,
                                         packed_bytes,
                                         "ds4_compressor_pool_kv") ||
        !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer,
                                         &g_compressor_pool_score_bytes,
                                         packed_bytes,
                                         "ds4_compressor_pool_score")) {
        return 0;
    }

    const uint64_t state_row_stride = (uint64_t)width * sizeof(float);
    const uint64_t pool_row_stride = (uint64_t)head_dim * sizeof(float);
    const NSUInteger curr_offset = (NSUInteger)4u * state_row_stride +
                                   (NSUInteger)head_dim * sizeof(float);
    if (!ds4_gpu_encode_concat_f32_dim1(cb,
                                          statekvbuf,
                                          ds4_gpu_tensor_offset(state_kv),
                                          4,
                                          state_row_stride,
                                          statekvbuf,
                                          ds4_gpu_tensor_offset(state_kv) + curr_offset,
                                          4,
                                          state_row_stride,
                                          g_compressor_pool_kv_buffer,
                                          0,
                                          head_dim,
                                          pool_row_stride) ||
        !ds4_gpu_encode_concat_f32_dim1(cb,
                                          statescbuf,
                                          ds4_gpu_tensor_offset(state_score),
                                          4,
                                          state_row_stride,
                                          statescbuf,
                                          ds4_gpu_tensor_offset(state_score) + curr_offset,
                                          4,
                                          state_row_stride,
                                          g_compressor_pool_score_buffer,
                                          0,
                                          head_dim,
                                          pool_row_stride)) {
        return 0;
    }

    return ds4_gpu_encode_dsv4_softmax_pool(cb,
                                              out,
                                              g_compressor_pool_kv_buffer,
                                              0,
                                              pool_row_stride,
                                              sizeof(float),
                                              packed_bytes,
                                              g_compressor_pool_score_buffer,
                                              0,
                                              pool_row_stride,
                                              sizeof(float),
                                              packed_bytes,
                                              8,
                                              head_dim,
                                              1);
}

static int ds4_gpu_encode_compressor_shift_ratio4(
        id<MTLCommandBuffer> cb,
        ds4_gpu_tensor    *state_kv,
        ds4_gpu_tensor    *state_score,
        uint32_t             width) {
    id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
    id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
    if (!cb || !statekvbuf || !statescbuf || !g_dsv4_ratio4_shift_pipeline || width == 0) return 0;

    ds4_gpu_dsv4_ratio4_shift_args args = { .width = width };
    const uint32_t n = 4u * width;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_dsv4_ratio4_shift_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:1];
    [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

int ds4_gpu_compressor_prefill_tensor(
        ds4_gpu_tensor       *comp_cache,
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        const ds4_gpu_tensor *kv,
        const ds4_gpu_tensor *sc,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint64_t                norm_offset,
        uint32_t                norm_type,
        uint32_t                head_dim,
        uint32_t                ratio,
        uint32_t                pos0,
        uint32_t                n_tokens,
        uint32_t                n_rot,
        uint32_t                n_ctx_orig,
        bool                    quantize_fp8,
        float                   freq_base,
        float                   freq_scale,
        float                   ext_factor,
        float                   attn_factor,
        float                   beta_fast,
        float                   beta_slow,
        float                   rms_eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
        head_dim == 0 || ratio == 0 || n_tokens == 0 ||
        n_rot > head_dim || (n_rot & 1u) != 0 ||
        (ape_type != 0u && ape_type != 1u) ||
        norm_type != 0u) {
        return 0;
    }

    @autoreleasepool {
        const uint32_t coff = ratio == 4u ? 2u : 1u;
        const uint32_t width = coff * head_dim;
        const uint32_t state_rows = coff * ratio;
        const uint32_t n_comp = n_tokens / ratio;
        const uint32_t cutoff = n_comp * ratio;
        const uint32_t rem = n_tokens - cutoff;
        const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
        const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
        const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
        const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
        const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
        const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);

        if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
            norm_offset > model_size || norm_bytes > model_size - norm_offset) {
            fprintf(stderr, "ds4: Metal compressor prefill tensor range is outside the mapped model\n");
            return 0;
        }

        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
        id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_cache);
        id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
        id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
        if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf ||
            ds4_gpu_tensor_bytes(kv) < kv_bytes ||
            ds4_gpu_tensor_bytes(sc) < kv_bytes ||
            ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
            ds4_gpu_tensor_bytes(state_score) < state_bytes ||
            (n_comp && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) {
            fprintf(stderr, "ds4: Metal compressor prefill received undersized buffers\n");
            return 0;
        }

        uint64_t ape_inner = 0;
        id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
        if (!apebuf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        int ok = 1;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) ok = 0;

        if (ok) {
            ok = ds4_gpu_encode_fill_f32_rows(cb,
                                                statekvbuf,
                                                ds4_gpu_tensor_offset(state_kv),
                                                width,
                                                state_rows,
                                                0.0f) &&
                 ds4_gpu_encode_fill_f32_rows(cb,
                                                statescbuf,
                                                ds4_gpu_tensor_offset(state_score),
                                                width,
                                                state_rows,
                                                ds4_gpu_negative_infinity());
        }

        if (ok && ratio == 4u) {
            int32_t rows_prev[4] = { 0, 1, 2, 3 };
            const int have_prev = cutoff >= ratio ? 1 : 0;
            const uint32_t prev_start = rem == 0 ? cutoff - ratio : cutoff - ratio;
            if (have_prev) {
                ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
                                                                     state_kv,
                                                                     state_score,
                                                                     kvbuf,
                                                                     ds4_gpu_tensor_offset(kv) +
                                                                             (NSUInteger)prev_start * width * sizeof(float),
                                                                     scbuf,
                                                                     ds4_gpu_tensor_offset(sc) +
                                                                             (NSUInteger)prev_start * width * sizeof(float),
                                                                     apebuf,
                                                                     (NSUInteger)ape_inner,
                                                                     ape_type,
                                                                     width,
                                                                     ratio,
                                                                     pos0 + prev_start,
                                                                     rows_prev,
                                                                     4,
                                                                     state_rows);
            }
            if (ok && rem != 0) {
                int32_t rows_cur[4];
                for (uint32_t i = 0; i < rem; i++) rows_cur[i] = (int32_t)(ratio + i);
                ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
                                                                     state_kv,
                                                                     state_score,
                                                                     kvbuf,
                                                                     ds4_gpu_tensor_offset(kv) +
                                                                             (NSUInteger)cutoff * width * sizeof(float),
                                                                     scbuf,
                                                                     ds4_gpu_tensor_offset(sc) +
                                                                             (NSUInteger)cutoff * width * sizeof(float),
                                                                     apebuf,
                                                                     (NSUInteger)ape_inner,
                                                                     ape_type,
                                                                     width,
                                                                     ratio,
                                                                     pos0 + cutoff,
                                                                     rows_cur,
                                                                     rem,
                                                                     state_rows);
            }
        } else if (ok && rem != 0) {
            int32_t rows[128];
            if (rem > (uint32_t)(sizeof(rows) / sizeof(rows[0]))) {
                fprintf(stderr, "ds4: Metal compressor prefill remainder exceeds local row list\n");
                ok = 0;
            } else {
                for (uint32_t i = 0; i < rem; i++) rows[i] = (int32_t)i;
                ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
                                                                     state_kv,
                                                                     state_score,
                                                                     kvbuf,
                                                                     ds4_gpu_tensor_offset(kv) +
                                                                             (NSUInteger)cutoff * width * sizeof(float),
                                                                     scbuf,
                                                                     ds4_gpu_tensor_offset(sc) +
                                                                             (NSUInteger)cutoff * width * sizeof(float),
                                                                     apebuf,
                                                                     (NSUInteger)ape_inner,
                                                                     ape_type,
                                                                     width,
                                                                     ratio,
                                                                     pos0 + cutoff,
                                                                     rows,
                                                                     rem,
                                                                     state_rows);
            }
        }

        if (ok && n_comp != 0) {
            const NSUInteger score_bytes = (NSUInteger)cutoff * width * sizeof(float);
            if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
                                                 &g_compressor_store_score_bytes,
                                                 score_bytes,
                                                 "ds4_compressor_store_score")) {
                ok = 0;
            }
            if (ok) {
                ok = ds4_gpu_encode_compressor_score_with_ape(cb,
                                                                 scbuf,
                                                                 ds4_gpu_tensor_offset(sc),
                                                                 g_compressor_store_score_buffer,
                                                                 0,
                                                                 apebuf,
                                                                 (NSUInteger)ape_inner,
                                                                 ape_type,
                                                                 width,
                                                                 ratio,
                                                                 pos0,
                                                                 cutoff);
            }

            if (ok && ratio == 4u) {
                const int direct_pool_mode =
                    ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp);
                if (ok && direct_pool_mode < 0) ok = 0;
                const uint64_t direct_output_bytes =
                    (uint64_t)n_comp * head_dim * sizeof(float);
                const uint64_t direct_input_bytes =
                    (uint64_t)n_comp * 4u * width * sizeof(float);
                const bool direct_pool_overlap = direct_pool_mode > 0 &&
                    ds4_gpu_buffer_ranges_overlap(
                        compbuf,
                        ds4_gpu_tensor_offset(comp_cache),
                        direct_output_bytes,
                        kvbuf,
                        ds4_gpu_tensor_offset(kv),
                        direct_input_bytes);
                const bool use_direct_pool =
                    direct_pool_mode > 0 && !direct_pool_overlap;
                const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float);
                if (ok && !use_direct_pool &&
                    (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer,
                                                      &g_compressor_pool_kv_bytes,
                                                      pack_bytes,
                                                      "ds4_compressor_pool_kv") ||
                     !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer,
                                                      &g_compressor_pool_score_bytes,
                                                      pack_bytes,
                                                      "ds4_compressor_pool_score"))) {
                    ok = 0;
                }
                const int pack_fusion_mode = use_direct_pool ? 0 :
                    ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim);
                if (ok && pack_fusion_mode < 0) ok = 0;
                const bool use_pack_fusion = pack_fusion_mode > 0;
                if (ok && use_direct_pool) {
                    ok = ds4_gpu_encode_compressor_ratio4_direct_pool(
                        cb,
                        comp_cache,
                        kvbuf,
                        ds4_gpu_tensor_offset(kv),
                        g_compressor_store_score_buffer,
                        0,
                        statekvbuf,
                        ds4_gpu_tensor_offset(state_kv),
                        statescbuf,
                        ds4_gpu_tensor_offset(state_score),
                        head_dim,
                        n_comp,
                        false);
                }
                if (ok && use_pack_fusion) {
                    ok = ds4_gpu_encode_compressor_pack_ratio4(
                        cb,
                        kvbuf,
                        ds4_gpu_tensor_offset(kv),
                        g_compressor_store_score_buffer,
                        0,
                        statekvbuf,
                        ds4_gpu_tensor_offset(state_kv),
                        statescbuf,
                        ds4_gpu_tensor_offset(state_score),
                        head_dim,
                        n_comp,
                        false);
                }
                if (ok && !use_direct_pool && !use_pack_fusion) {
                    ok = ds4_gpu_encode_fill_f32_rows(cb,
                                                        g_compressor_pool_kv_buffer,
                                                        0,
                                                        head_dim,
                                                        8u * n_comp,
                                                        0.0f) &&
                         ds4_gpu_encode_fill_f32_rows(cb,
                                                        g_compressor_pool_score_buffer,
                                                        0,
                                                        head_dim,
                                                        8u * n_comp,
                                                        ds4_gpu_negative_infinity());
                }
                if (ok && !use_direct_pool && !use_pack_fusion) {
                    const uint64_t src_row_stride = (uint64_t)width * sizeof(float);
                    const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride;
                    const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float);
                    const uint64_t dst_plane_stride = 8ull * dst_row_stride;
                    ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                         kvbuf,
                                                         ds4_gpu_tensor_offset(kv) +
                                                                 (NSUInteger)head_dim * sizeof(float),
                                                         g_compressor_pool_kv_buffer,
                                                         (NSUInteger)4u * head_dim * sizeof(float),
                                                         head_dim,
                                                         ratio,
                                                         n_comp,
                                                         src_row_stride,
                                                         src_plane_stride,
                                                         dst_row_stride,
                                                         dst_plane_stride) &&
                         ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                         g_compressor_store_score_buffer,
                                                         (NSUInteger)head_dim * sizeof(float),
                                                         g_compressor_pool_score_buffer,
                                                         (NSUInteger)4u * head_dim * sizeof(float),
                                                         head_dim,
                                                         ratio,
                                                         n_comp,
                                                         src_row_stride,
                                                         src_plane_stride,
                                                         dst_row_stride,
                                                         dst_plane_stride);
                }
                if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) {
                    const uint64_t src_row_stride = (uint64_t)width * sizeof(float);
                    const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride;
                    const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float);
                    const uint64_t dst_plane_stride = 8ull * dst_row_stride;
                    ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                         kvbuf,
                                                         ds4_gpu_tensor_offset(kv),
                                                         g_compressor_pool_kv_buffer,
                                                         dst_plane_stride,
                                                         head_dim,
                                                         ratio,
                                                         n_comp - 1u,
                                                         src_row_stride,
                                                         src_plane_stride,
                                                         dst_row_stride,
                                                         dst_plane_stride) &&
                         ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                         g_compressor_store_score_buffer,
                                                         0,
                                                         g_compressor_pool_score_buffer,
                                                         dst_plane_stride,
                                                         head_dim,
                                                         ratio,
                                                         n_comp - 1u,
                                                         src_row_stride,
                                                         src_plane_stride,
                                                         dst_row_stride,
                                                         dst_plane_stride);
                }
                if (ok && !use_direct_pool) {
                    ok = ds4_gpu_encode_dsv4_softmax_pool(cb,
                                                            comp_cache,
                                                            g_compressor_pool_kv_buffer,
                                                            0,
                                                            (uint64_t)head_dim * sizeof(float),
                                                            sizeof(float),
                                                            8ull * head_dim * sizeof(float),
                                                            g_compressor_pool_score_buffer,
                                                            0,
                                                            (uint64_t)head_dim * sizeof(float),
                                                            sizeof(float),
                                                            8ull * head_dim * sizeof(float),
                                                            8,
                                                            head_dim,
                                                            n_comp);
                }
            } else if (ok) {
                const uint64_t row_stride = (uint64_t)width * sizeof(float);
                ok = ds4_gpu_encode_dsv4_softmax_pool(cb,
                                                        comp_cache,
                                                        kvbuf,
                                                        ds4_gpu_tensor_offset(kv),
                                                        row_stride,
                                                        sizeof(float),
                                                        (uint64_t)ratio * row_stride,
                                                        g_compressor_store_score_buffer,
                                                        0,
                                                        row_stride,
                                                        sizeof(float),
                                                        (uint64_t)ratio * row_stride,
                                                        ratio,
                                                        head_dim,
                                                        n_comp);
            }
        }

        if (ok && n_comp != 0) {
            ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache,
                                                       comp_cache,
                                                       model_map,
                                                       model_size,
                                                       norm_offset,
                                                       head_dim,
                                                       n_comp,
                                                       rms_eps) != 0;
        }
        if (ok && n_comp != 0 && n_rot != 0) {
            ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args(
                n_comp, 1, head_dim, n_rot, n_ctx_orig, false,
                freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
            cb = ds4_gpu_command_buffer(&owned);
            ok = cb && !owned &&
                 ds4_gpu_encode_rope_tail_inplace(cb,
                                                    compbuf,
                                                    ds4_gpu_tensor_offset(comp_cache),
                                                    &rope_args,
                                                    n_comp,
                                                    1,
                                                    head_dim,
                                                    pos0,
                                                    ratio);
        }
        if (ok && n_comp != 0 && quantize_fp8) {
            ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0;
        }

        if (!had_batch) {
            const int end_ok = ds4_gpu_end_commands();
            ok = end_ok && ok;
        }
        return ok ? 1 : 0;
    }
}

int ds4_gpu_compressor_prefill_ratio4_replay_tensor(
        ds4_gpu_tensor       *comp_cache,
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        const ds4_gpu_tensor *kv,
        const ds4_gpu_tensor *sc,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint64_t                norm_offset,
        uint32_t                norm_type,
        uint32_t                head_dim,
        uint32_t                pos0,
        uint32_t                n_tokens,
        uint32_t                n_rot,
        uint32_t                n_ctx_orig,
        bool                    quantize_fp8,
        float                   freq_base,
        float                   freq_scale,
        float                   ext_factor,
        float                   attn_factor,
        float                   beta_fast,
        float                   beta_slow,
        float                   rms_eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
        head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 ||
        n_rot > head_dim || (n_rot & 1u) != 0 ||
        (ape_type != 0u && ape_type != 1u) ||
        norm_type != 0u) {
        return 0;
    }

    @autoreleasepool {
        const uint32_t ratio = 4u;
        const uint32_t width = 2u * head_dim;
        const uint32_t state_rows = 8u;
        const uint32_t n_comp = n_tokens / ratio;
        const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
        const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
        const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
        const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
        const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
        const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);

        if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
            norm_offset > model_size || norm_bytes > model_size - norm_offset) {
            fprintf(stderr, "ds4: Metal compressor replay tensor range is outside the mapped model\n");
            return 0;
        }

        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
        id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_cache);
        id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
        id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
        if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf ||
            ds4_gpu_tensor_bytes(kv) < kv_bytes ||
            ds4_gpu_tensor_bytes(sc) < kv_bytes ||
            ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
            ds4_gpu_tensor_bytes(state_score) < state_bytes ||
            ds4_gpu_tensor_bytes(comp_cache) < comp_bytes) {
            fprintf(stderr, "ds4: Metal compressor replay received undersized buffers\n");
            return 0;
        }

        uint64_t ape_inner = 0;
        id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
        if (!apebuf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        int ok = 1;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) ok = 0;

        const NSUInteger score_bytes = (NSUInteger)n_tokens * width * sizeof(float);
        const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float);
        const int direct_pool_mode =
            ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp);
        if (ok && direct_pool_mode < 0) ok = 0;
        const uint64_t direct_output_bytes =
            (uint64_t)n_comp * head_dim * sizeof(float);
        const uint64_t direct_input_bytes =
            (uint64_t)n_tokens * width * sizeof(float);
        const uint64_t direct_state_bytes =
            (uint64_t)4u * width * sizeof(float);
        const bool direct_pool_overlap = direct_pool_mode > 0 &&
            (ds4_gpu_buffer_ranges_overlap(
                 compbuf,
                 ds4_gpu_tensor_offset(comp_cache),
                 direct_output_bytes,
                 kvbuf,
                 ds4_gpu_tensor_offset(kv),
                 direct_input_bytes) ||
             ds4_gpu_buffer_ranges_overlap(
                 compbuf,
                 ds4_gpu_tensor_offset(comp_cache),
                 direct_output_bytes,
                 statekvbuf,
                 ds4_gpu_tensor_offset(state_kv),
                 direct_state_bytes) ||
             ds4_gpu_buffer_ranges_overlap(
                 compbuf,
                 ds4_gpu_tensor_offset(comp_cache),
                 direct_output_bytes,
                 statescbuf,
                 ds4_gpu_tensor_offset(state_score),
                 direct_state_bytes));
        const bool use_direct_pool =
            direct_pool_mode > 0 && !direct_pool_overlap;
        if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
                                                    &g_compressor_store_score_bytes,
                                                    score_bytes,
                                                    "ds4_compressor_store_score") ||
                   (!use_direct_pool &&
                    (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer,
                                                     &g_compressor_pool_kv_bytes,
                                                     pack_bytes,
                                                     "ds4_compressor_pool_kv") ||
                     !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer,
                                                     &g_compressor_pool_score_bytes,
                                                     pack_bytes,
                                                     "ds4_compressor_pool_score"))))) {
            ok = 0;
        }

        if (ok) {
            ok = ds4_gpu_encode_compressor_score_with_ape(cb,
                                                            scbuf,
                                                            ds4_gpu_tensor_offset(sc),
                                                            g_compressor_store_score_buffer,
                                                            0,
                                                            apebuf,
                                                            (NSUInteger)ape_inner,
                                                            ape_type,
                                                            width,
                                                            ratio,
                                                            pos0,
                                                            n_tokens);
        }

        const uint64_t src_row_stride = (uint64_t)width * sizeof(float);
        const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride;
        const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float);
        const uint64_t dst_plane_stride = 8ull * dst_row_stride;
        const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv);
        const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score);
        const int pack_fusion_mode = use_direct_pool ? 0 :
            ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim);
        if (ok && pack_fusion_mode < 0) ok = 0;
        const bool use_pack_fusion = pack_fusion_mode > 0;

        if (ok && use_direct_pool) {
            ok = ds4_gpu_encode_compressor_ratio4_direct_pool(
                cb,
                comp_cache,
                kvbuf,
                ds4_gpu_tensor_offset(kv),
                g_compressor_store_score_buffer,
                0,
                statekvbuf,
                state_off,
                statescbuf,
                state_score_off,
                head_dim,
                n_comp,
                true);
        }

        if (ok && use_pack_fusion) {
            ok = ds4_gpu_encode_compressor_pack_ratio4(
                cb,
                kvbuf,
                ds4_gpu_tensor_offset(kv),
                g_compressor_store_score_buffer,
                0,
                statekvbuf,
                state_off,
                statescbuf,
                state_score_off,
                head_dim,
                n_comp,
                true);
        }

        if (ok && !use_direct_pool && !use_pack_fusion) {
            ok = ds4_gpu_encode_fill_f32_rows(cb,
                                                g_compressor_pool_kv_buffer,
                                                0,
                                                head_dim,
                                                8u * n_comp,
                                                0.0f) &&
                 ds4_gpu_encode_fill_f32_rows(cb,
                                                g_compressor_pool_score_buffer,
                                                0,
                                                head_dim,
                                                8u * n_comp,
                                                ds4_gpu_negative_infinity());
        }

        if (ok && !use_direct_pool && !use_pack_fusion) {
            /*
             * The aligned nonzero ratio-4 path replays the current ubatch
             * compressor, but seeds the first compressed row with the previous
             * compressor state. Rows 0..3 are the previous half, rows 4..7 are
             * the current half.
             */
            ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                 statekvbuf,
                                                 state_off,
                                                 g_compressor_pool_kv_buffer,
                                                 0,
                                                 head_dim,
                                                 ratio,
                                                 1,
                                                 src_row_stride,
                                                 (uint64_t)ratio * src_row_stride,
                                                 dst_row_stride,
                                                 dst_plane_stride) &&
                 ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                 statescbuf,
                                                 state_score_off,
                                                 g_compressor_pool_score_buffer,
                                                 0,
                                                 head_dim,
                                                 ratio,
                                                 1,
                                                 src_row_stride,
                                                 (uint64_t)ratio * src_row_stride,
                                                 dst_row_stride,
                                                 dst_plane_stride);
        }
        if (ok && !use_direct_pool && !use_pack_fusion) {
            ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                 kvbuf,
                                                 ds4_gpu_tensor_offset(kv) +
                                                         (NSUInteger)head_dim * sizeof(float),
                                                 g_compressor_pool_kv_buffer,
                                                 (NSUInteger)4u * head_dim * sizeof(float),
                                                 head_dim,
                                                 ratio,
                                                 n_comp,
                                                 src_row_stride,
                                                 src_plane_stride,
                                                 dst_row_stride,
                                                 dst_plane_stride) &&
                 ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                 g_compressor_store_score_buffer,
                                                 (NSUInteger)head_dim * sizeof(float),
                                                 g_compressor_pool_score_buffer,
                                                 (NSUInteger)4u * head_dim * sizeof(float),
                                                 head_dim,
                                                 ratio,
                                                 n_comp,
                                                 src_row_stride,
                                                 src_plane_stride,
                                                 dst_row_stride,
                                                 dst_plane_stride);
        }
        if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) {
            ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                 kvbuf,
                                                 ds4_gpu_tensor_offset(kv),
                                                 g_compressor_pool_kv_buffer,
                                                 dst_plane_stride,
                                                 head_dim,
                                                 ratio,
                                                 n_comp - 1u,
                                                 src_row_stride,
                                                 src_plane_stride,
                                                 dst_row_stride,
                                                 dst_plane_stride) &&
                 ds4_gpu_encode_cpy_f32_f32_3d(cb,
                                                 g_compressor_store_score_buffer,
                                                 0,
                                                 g_compressor_pool_score_buffer,
                                                 dst_plane_stride,
                                                 head_dim,
                                                 ratio,
                                                 n_comp - 1u,
                                                 src_row_stride,
                                                 src_plane_stride,
                                                 dst_row_stride,
                                                 dst_plane_stride);
        }
        if (ok && !use_direct_pool) {
            ok = ds4_gpu_encode_dsv4_softmax_pool(cb,
                                                    comp_cache,
                                                    g_compressor_pool_kv_buffer,
                                                    0,
                                                    dst_row_stride,
                                                    sizeof(float),
                                                    dst_plane_stride,
                                                    g_compressor_pool_score_buffer,
                                                    0,
                                                    dst_row_stride,
                                                    sizeof(float),
                                                    dst_plane_stride,
                                                    8,
                                                    head_dim,
                                                    n_comp);
        }
        if (ok) {
            ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache,
                                                       comp_cache,
                                                       model_map,
                                                       model_size,
                                                       norm_offset,
                                                       head_dim,
                                                       n_comp,
                                                       rms_eps) != 0;
        }
        if (ok && n_rot != 0) {
            ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args(
                n_comp, 1, head_dim, n_rot, n_ctx_orig, false,
                freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
            cb = ds4_gpu_command_buffer(&owned);
            ok = cb && !owned &&
                 ds4_gpu_encode_rope_tail_inplace(cb,
                                                    compbuf,
                                                    ds4_gpu_tensor_offset(comp_cache),
                                                    &rope_args,
                                                    n_comp,
                                                    1,
                                                    head_dim,
                                                    pos0,
                                                    ratio);
        }
        if (ok && quantize_fp8) {
            ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0;
        }

        if (ok) {
            ok = ds4_gpu_encode_fill_f32_rows(cb,
                                                statekvbuf,
                                                state_off,
                                                width,
                                                state_rows,
                                                0.0f) &&
                 ds4_gpu_encode_fill_f32_rows(cb,
                                                statescbuf,
                                                state_score_off,
                                                width,
                                                state_rows,
                                                ds4_gpu_negative_infinity());
        }
        if (ok) {
            int32_t rows_prev[4] = { 0, 1, 2, 3 };
            const uint32_t prev_start = n_tokens - ratio;
            ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
                                                                 state_kv,
                                                                 state_score,
                                                                 kvbuf,
                                                                 ds4_gpu_tensor_offset(kv) +
                                                                         (NSUInteger)prev_start * width * sizeof(float),
                                                                 scbuf,
                                                                 ds4_gpu_tensor_offset(sc) +
                                                                         (NSUInteger)prev_start * width * sizeof(float),
                                                                 apebuf,
                                                                 (NSUInteger)ape_inner,
                                                                 ape_type,
                                                                 width,
                                                                 ratio,
                                                                 pos0 + prev_start,
                                                                 rows_prev,
                                                                 ratio,
                                                                 state_rows);
        }

        if (!had_batch) {
            const int end_ok = ds4_gpu_end_commands();
            ok = end_ok && ok;
        }
        return ok ? 1 : 0;
    }
}

int ds4_gpu_compressor_prefill_state_ratio4_tensor(
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        const ds4_gpu_tensor *kv_tail,
        const ds4_gpu_tensor *sc_tail,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint32_t                head_dim,
        uint32_t                pos0) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map ||
        head_dim == 0 || (ape_type != 0u && ape_type != 1u)) {
        return 0;
    }

    @autoreleasepool {
        const uint32_t ratio = 4u;
        const uint32_t width = 2u * head_dim;
        const uint32_t state_rows = 8u;
        const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
        const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float);
        const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
        const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape;

        if (ape_offset > model_size || ape_bytes > model_size - ape_offset) {
            fprintf(stderr, "ds4: Metal compressor prefill-state APE range is outside the mapped model\n");
            return 0;
        }

        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv_tail);
        id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc_tail);
        id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
        id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
        if (!kvbuf || !scbuf || !statekvbuf || !statescbuf ||
            ds4_gpu_tensor_bytes(kv_tail) < tail_bytes ||
            ds4_gpu_tensor_bytes(sc_tail) < tail_bytes ||
            ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
            ds4_gpu_tensor_bytes(state_score) < state_bytes) {
            fprintf(stderr, "ds4: Metal compressor prefill-state received undersized buffers\n");
            return 0;
        }

        uint64_t ape_inner = 0;
        id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
        if (!apebuf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        int ok = 1;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) ok = 0;

        if (ok) {
            ok = ds4_gpu_encode_fill_f32_rows(cb,
                                                statekvbuf,
                                                ds4_gpu_tensor_offset(state_kv),
                                                width,
                                                state_rows,
                                                0.0f) &&
                 ds4_gpu_encode_fill_f32_rows(cb,
                                                statescbuf,
                                                ds4_gpu_tensor_offset(state_score),
                                                width,
                                                state_rows,
                                                ds4_gpu_negative_infinity());
        }
        if (ok) {
            int32_t rows[4] = { 0, 1, 2, 3 };
            ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
                                                                 state_kv,
                                                                 state_score,
                                                                 kvbuf,
                                                                 ds4_gpu_tensor_offset(kv_tail),
                                                                 scbuf,
                                                                 ds4_gpu_tensor_offset(sc_tail),
                                                                 apebuf,
                                                                 (NSUInteger)ape_inner,
                                                                 ape_type,
                                                                 width,
                                                                 ratio,
                                                                 pos0,
                                                                 rows,
                                                                 ratio,
                                                                 state_rows);
        }

        if (!had_batch) {
            const int end_ok = ds4_gpu_end_commands();
            ok = end_ok && ok;
        }
        return ok ? 1 : 0;
    }
}

int ds4_gpu_compressor_update_tensor(
        const ds4_gpu_tensor *kv_cur,
        const ds4_gpu_tensor *sc_cur,
        ds4_gpu_tensor       *state_kv,
        ds4_gpu_tensor       *state_score,
        ds4_gpu_tensor       *comp_cache,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                ape_offset,
        uint32_t                ape_type,
        uint64_t                norm_offset,
        uint32_t                norm_type,
        uint32_t                head_dim,
        uint32_t                ratio,
        uint32_t                pos,
        uint32_t                comp_row,
        uint32_t                n_rot,
        uint32_t                n_ctx_orig,
        float                   freq_base,
        float                   freq_scale,
        float                   ext_factor,
        float                   attn_factor,
        float                   beta_fast,
        float                   beta_slow,
        float                   rms_eps,
        bool                    state_already_stored) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
        !model_map || head_dim == 0 || ratio == 0 ||
        n_rot > head_dim || (n_rot & 1u) != 0 ||
        (ape_type != 0u && ape_type != 1u) ||
        norm_type != 0u) {
        return 0;
    }

    @autoreleasepool {
        const uint32_t coff = ratio == 4u ? 2u : 1u;
        const uint32_t width = coff * head_dim;
        const uint32_t state_rows = coff * ratio;
        const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u;
        const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
        const uint64_t kv_bytes = (uint64_t)width * sizeof(float);
        const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
        const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float);
        const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
        const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);

        if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
            norm_offset > model_size || norm_bytes > model_size - norm_offset) {
            fprintf(stderr, "ds4: Metal compressor tensor range is outside the mapped model\n");
            return 0;
        }

        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv_cur);
        id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc_cur);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_cache);
        if (!kvbuf || !scbuf || !compbuf ||
            ds4_gpu_tensor_bytes(kv_cur) < kv_bytes ||
            ds4_gpu_tensor_bytes(sc_cur) < kv_bytes ||
            ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
            ds4_gpu_tensor_bytes(state_score) < state_bytes ||
            (emit && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) {
            fprintf(stderr, "ds4: Metal compressor update received undersized buffers\n");
            return 0;
        }

        if (!state_already_stored) {
            const bool use_store_one =
                getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL;
            const int store_ok = use_store_one
                ? ds4_gpu_compressor_store_one_tensor(kv_cur,
                                                        sc_cur,
                                                        state_kv,
                                                        state_score,
                                                        model_map,
                                                        model_size,
                                                        ape_offset,
                                                        ape_type,
                                                        width,
                                                        ratio,
                                                        pos)
                : ds4_gpu_compressor_store_batch_tensor(kv_cur,
                                                          sc_cur,
                                                          state_kv,
                                                          state_score,
                                                          model_map,
                                                          model_size,
                                                          ape_offset,
                                                          ape_type,
                                                          head_dim,
                                                          ratio,
                                                          pos,
                                                          1);
            if (!store_ok) {
                return 0;
            }
        }
        if (!emit) return 1;

        ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
                comp_cache,
                (uint64_t)comp_row * head_dim * sizeof(float),
                (uint64_t)head_dim * sizeof(float));
        if (!comp_row_view) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        int ok = cb &&
                 ds4_gpu_encode_compressor_pool(cb,
                                                  comp_row_view,
                                                  state_kv,
                                                  state_score,
                                                  head_dim,
                                                  ratio);
        if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor DS4 softmax pool");
        if (ok) {
            ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view,
                                                       comp_row_view,
                                                       model_map,
                                                       model_size,
                                                       norm_offset,
                                                       head_dim,
                                                       1,
                                                       rms_eps) != 0;
        }
        if (ok) {
            const uint32_t comp_pos = pos + 1u - ratio;
            ok = ds4_gpu_rope_tail_tensor(comp_row_view,
                                            1,
                                            1,
                                            head_dim,
                                            n_rot,
                                            comp_pos,
                                            n_ctx_orig,
                                            false,
                                            freq_base,
                                            freq_scale,
                                            ext_factor,
                                            attn_factor,
                                            beta_fast,
                                            beta_slow) != 0;
        }
        if (ok && ratio == 4u) {
            cb = ds4_gpu_command_buffer(&owned);
            ok = cb &&
                 ds4_gpu_encode_compressor_shift_ratio4(cb,
                                                          state_kv,
                                                          state_score,
                                                          width);
            if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor ratio4 state shift");
        }
        ds4_gpu_tensor_free(comp_row_view);
        if (!ok) return 0;
    }

    return 1;
}

static int ds4_gpu_encode_fill_f32_rows(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        buf,
        NSUInteger           offset,
        uint32_t             width,
        uint32_t             rows,
        float                value) {
    if (!cb || !buf || width == 0 || rows == 0 || (width & 3u) != 0) return 0;

    ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, 1, 0.0f, 0.0f);
    args.val = value;

    NSUInteger nth_max = g_unary_fill_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth_max > 256u) nth_max = 256u;
    NSUInteger nth = (NSUInteger)args.ne00;
    if (nth > nth_max) nth = nth_max;
    if (nth == 0) nth = 1u;
    const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_unary_fill_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:buf offset:offset atIndex:1];
    [enc setBuffer:buf offset:offset atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01,
                                          (NSUInteger)args.ne02,
                                          (NSUInteger)args.ne03)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

int ds4_gpu_attention_output_q8_batch_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *low,
        ds4_gpu_tensor       *group_tmp,
        ds4_gpu_tensor       *low_tmp,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                out_a_offset,
        uint64_t                out_b_offset,
        uint64_t                group_dim,
        uint64_t                rank,
        uint32_t                n_groups,
        uint64_t                out_dim,
        const ds4_gpu_tensor *heads,
        uint32_t                n_tokens) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !low || !group_tmp || !low_tmp || !heads || !model_map ||
        group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 ||
        group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t low_dim = (uint64_t)n_groups * rank;
        if ((group_dim % 32u) != 0 || (low_dim % 32u) != 0 || low_dim > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal attention output batch received invalid q8 dimensions\n");
            return 0;
        }
        const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
        const uint64_t row_b_bytes = (low_dim / 32u) * 34u;
        const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
        const uint64_t out_b_bytes = out_dim * row_b_bytes;
        if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset ||
            out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) {
            fprintf(stderr, "ds4: Metal attention output batch weights are outside the mapped model\n");
            return 0;
        }

        const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float);
        const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float);
        if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(low) < low_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal attention output batch received undersized buffers\n");
            return 0;
        }
        (void)group_tmp;
        (void)low_tmp;

        const bool use_direct_low =
            n_tokens < 32u && getenv("DS4_METAL_DISABLE_ATTN_OUT_LOW_DIRECT") == NULL;
        /* The exported TensorOps attention-output kernel is a 64-token tile.
         * Keep this on full tiles only; smaller multiples of 32 use the legacy
         * path instead of relying on cooperative tensor partial RHS bounds. */
        const bool use_mpp_low =
            n_tokens >= 32u &&
            (n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 &&
            ds4_gpu_use_mpp_attn_out_low_matmul();
        const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t);
        id<MTLBuffer> group_ids_buffer = nil;
        if (!use_direct_low && !use_mpp_low) {
            if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) {
                group_ids_buffer =
                    ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids");
                if (!group_ids_buffer) {
                    return 0;
                }
            } else {
                if (!ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer,
                                                     &g_attn_out_group_ids_bytes,
                                                     ids_bytes,
                                                     "ds4_attention_output_group_ids")) {
                    return 0;
                }
                group_ids_buffer = g_attn_out_group_ids_buffer;
            }
            int32_t *ids = (int32_t *)[group_ids_buffer contents];
            for (uint32_t t = 0; t < n_tokens; t++) {
                for (uint32_t group = 0; group < n_groups; group++) {
                    ids[(uint64_t)t * n_groups + group] = (int32_t)group;
                }
            }
        }

        uint64_t out_a_inner = 0;
        id<MTLBuffer> out_a_buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                       out_a_offset, out_a_bytes,
                                       &out_a_inner);
        if (!out_a_buf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        bool ok = true;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) {
            ok = false;
        }
        const bool attn_out_profile =
            getenv("DS4_METAL_ATTN_OUT_STAGE_PROFILE") != NULL && g_batch_cb != nil;
        if (ok && attn_out_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                ok = false;
            } else {
                cb = ds4_gpu_command_buffer(&owned);
                if (!cb || owned) ok = false;
            }
        }
        double attn_out_t0 = attn_out_profile ? ds4_gpu_now_ms() : 0.0;
#define DS4_METAL_PROFILE_ATTN_OUT_STAGE(name) do { \
            if (ok && attn_out_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = false; \
                } else { \
                    const double now_ms = ds4_gpu_now_ms(); \
                    fprintf(stderr, \
                            "ds4: Metal attention output stage tokens=%u %s=%.3f ms\n", \
                            n_tokens, (name), now_ms - attn_out_t0); \
                    attn_out_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = false; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb || owned) ok = false; \
                    } \
                } \
            } \
        } while (0)

        if (ok) {
            /*
             * Batched attention-output projections switch from the vector
             * kernel to the SIMD matrix kernel once the batch has at least 32
             * tokens.  This preserves the single-token generation path while
             * keeping prefill accumulation stable.
             */
            if (use_mpp_low) {
                ds4_gpu_mul_mm_id_args mm_args =
                    ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim,
                                                  (uint32_t)rank,
                                                  n_groups,
                                                  row_a_bytes,
                                                  (uint64_t)rank * row_a_bytes,
                                                  n_groups,
                                                  n_groups,
                                                  n_tokens);
                /*
                 * Direct RHS lets MPP read the dense low-rank activation tile
                 * directly from device memory instead of staging a second
                 * threadgroup tile.  The retained attention-output path is the
                 * 64-token direct-RHS kernel; the older staged-RHS and 32-token
                 * variants were not kept as alternate runtime modes.
                 */
                const char *attn_out_pipeline_name =
                    "kernel_attn_out_low_q8_0_mpp_direct_rhs_n64";
                id<MTLComputePipelineState> mm_pipeline =
                    ds4_gpu_get_mul_mm_id_pipeline(attn_out_pipeline_name, false);
                ok = ds4_gpu_encode_attn_out_low_q8_mpp(cb,
                                                          mm_pipeline,
                                                          &mm_args,
                                                          out_a_buf,
                                                          (NSUInteger)out_a_inner,
                                                          ds4_gpu_tensor_buffer(heads),
                                                          ds4_gpu_tensor_offset(heads),
                                                          ds4_gpu_tensor_buffer(low),
                                                          ds4_gpu_tensor_offset(low)) != 0;
                if (!ok) {
                    ds4_gpu_warn_mpp_fallback();
                    if (ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) {
                        if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) {
                            group_ids_buffer =
                                ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids");
                        } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer,
                                                                   &g_attn_out_group_ids_bytes,
                                                                   ids_bytes,
                                                                   "ds4_attention_output_group_ids")) {
                            group_ids_buffer = g_attn_out_group_ids_buffer;
                        }
                        if (group_ids_buffer) {
                            int32_t *ids = (int32_t *)[group_ids_buffer contents];
                            for (uint32_t t = 0; t < n_tokens; t++) {
                                for (uint32_t group = 0; group < n_groups; group++) {
                                    ids[(uint64_t)t * n_groups + group] = (int32_t)group;
                                }
                            }
                            ds4_gpu_mul_mm_id_map_args map_args =
                                ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim,
                                                                  n_groups,
                                                                  n_groups,
                                                                  n_groups,
                                                                  n_tokens);
                            id<MTLComputePipelineState> map_pipeline =
                                ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups));
                            id<MTLComputePipelineState> fallback_pipeline =
                                ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false);
                            ok = ds4_gpu_encode_mul_mm_id(cb,
                                                            map_pipeline,
                                                            fallback_pipeline,
                                                            &map_args,
                                                            &mm_args,
                                                            out_a_buf,
                                                            (NSUInteger)out_a_inner,
                                                            ds4_gpu_tensor_buffer(heads),
                                                            ds4_gpu_tensor_offset(heads),
                                                            ds4_gpu_tensor_buffer(low),
                                                            ds4_gpu_tensor_offset(low),
                                                            group_ids_buffer,
                                                            0) != 0;
                        }
                    }
                }
            } else if (n_tokens >= 32u && ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) {
                ds4_gpu_mul_mm_id_map_args map_args =
                    ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim,
                                                      n_groups,
                                                      n_groups,
                                                      n_groups,
                                                      n_tokens);
                ds4_gpu_mul_mm_id_args mm_args =
                    ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim,
                                                  (uint32_t)rank,
                                                  n_groups,
                                                  row_a_bytes,
                                                  (uint64_t)rank * row_a_bytes,
                                                  n_groups,
                                                  n_groups,
                                                  n_tokens);
                id<MTLComputePipelineState> map_pipeline =
                    ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups));
                id<MTLComputePipelineState> mm_pipeline =
                    ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false);
                ok = ds4_gpu_encode_mul_mm_id(cb,
                                                map_pipeline,
                                                mm_pipeline,
                                                &map_args,
                                                &mm_args,
                                                out_a_buf,
                                                (NSUInteger)out_a_inner,
                                                ds4_gpu_tensor_buffer(heads),
                                                ds4_gpu_tensor_offset(heads),
                                                ds4_gpu_tensor_buffer(low),
                                                ds4_gpu_tensor_offset(low),
                                                group_ids_buffer,
                                                0) != 0;
            } else if (use_direct_low) {
                ds4_gpu_mul_mv_id_args args = {
                    .nei0 = (int32_t)n_groups,
                    .nei1 = (int32_t)n_tokens,
                    .nbi1 = 0,
                    .ne00 = (int32_t)group_dim,
                    .ne01 = (int32_t)rank,
                    .ne02 = (int32_t)n_groups,
                    .nb00 = 34,
                    .nb01 = row_a_bytes,
                    .nb02 = (uint64_t)rank * row_a_bytes,
                    .ne10 = (int32_t)group_dim,
                    .ne11 = (int32_t)n_groups,
                    .ne12 = (int32_t)n_tokens,
                    .ne13 = 1,
                    .nb10 = sizeof(float),
                    .nb11 = (uint64_t)group_dim * sizeof(float),
                    .nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
                    .ne0 = (int32_t)rank,
                    .ne1 = (int32_t)n_groups,
                    .nb1 = (uint64_t)rank * sizeof(float),
                    .nr0 = 2,
                };
                id<MTLComputePipelineState> pipeline =
                    ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4);
                ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
                                                             pipeline,
                                                             &args,
                                                             out_a_buf,
                                                             (NSUInteger)out_a_inner,
                                                             ds4_gpu_tensor_buffer(heads),
                                                             ds4_gpu_tensor_offset(heads),
                                                             ds4_gpu_tensor_buffer(low),
                                                             ds4_gpu_tensor_offset(low),
                                                             32u * 2u * sizeof(float),
                                                             4,
                                                             true) != 0;
            } else {
                ds4_gpu_mul_mv_id_args args = {
                    .nei0 = (int32_t)n_groups,
                    .nei1 = (int32_t)n_tokens,
                    .nbi1 = (uint64_t)n_groups * sizeof(int32_t),
                    .ne00 = (int32_t)group_dim,
                    .ne01 = (int32_t)rank,
                    .ne02 = (int32_t)n_groups,
                    .nb00 = 34,
                    .nb01 = row_a_bytes,
                    .nb02 = (uint64_t)rank * row_a_bytes,
                    .ne10 = (int32_t)group_dim,
                    .ne11 = (int32_t)n_groups,
                    .ne12 = (int32_t)n_tokens,
                    .ne13 = 1,
                    .nb10 = sizeof(float),
                    .nb11 = (uint64_t)group_dim * sizeof(float),
                    .nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
                    .ne0 = (int32_t)rank,
                    .ne1 = (int32_t)n_groups,
                    .nb1 = (uint64_t)rank * sizeof(float),
                    .nr0 = 2,
                };
                id<MTLComputePipelineState> pipeline =
                    ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4);
                ok = ds4_gpu_encode_mul_mv_id(cb,
                                                pipeline,
                                                &args,
                                                out_a_buf,
                                                (NSUInteger)out_a_inner,
                                                ds4_gpu_tensor_buffer(heads),
                                                ds4_gpu_tensor_offset(heads),
                                                ds4_gpu_tensor_buffer(low),
                                                ds4_gpu_tensor_offset(low),
                                                group_ids_buffer,
                                                0,
                                                32u * 2u * sizeof(float),
                                                4,
                                                true) != 0;
            }
        }
        DS4_METAL_PROFILE_ATTN_OUT_STAGE("low_proj");

        if (ok) {
            ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size,
                                              out_b_offset,
                                              low_dim, out_dim, low, n_tokens) != 0;
        }
        DS4_METAL_PROFILE_ATTN_OUT_STAGE("out_proj");

        if (!had_batch) {
            ok = ds4_gpu_end_commands() != 0 && ok;
        }
#undef DS4_METAL_PROFILE_ATTN_OUT_STAGE
        return ok ? 1 : 0;
    }
}

int ds4_gpu_attention_output_q4_K_batch_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *low,
        ds4_gpu_tensor       *group_tmp,
        ds4_gpu_tensor       *low_tmp,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                out_a_offset,
        uint64_t                out_b_offset,
        uint32_t                out_b_type,
        uint64_t                group_dim,
        uint64_t                rank,
        uint32_t                n_groups,
        uint64_t                out_dim,
        const ds4_gpu_tensor *heads,
        uint32_t                n_tokens) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !low || !heads || !model_map ||
        group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 ||
        group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) {
        return 0;
    }
    if (n_tokens < 32u) return 0;

    @autoreleasepool {
        const uint64_t low_dim = (uint64_t)n_groups * rank;
        if ((group_dim % 256u) != 0 || (low_dim % 256u) != 0 || low_dim > UINT32_MAX) {
            return 0;
        }

        uint64_t row_a_bytes = 0;
        uint64_t row_b_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, (uint32_t)group_dim, &row_a_bytes) ||
            !ds4_gpu_quant_row_bytes(out_b_type, (uint32_t)low_dim, &row_b_bytes)) {
            return 0;
        }

        const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
        const uint64_t out_b_bytes = out_dim * row_b_bytes;
        if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset ||
            out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) {
            fprintf(stderr, "ds4: Metal Q4 attention output batch weights are outside the mapped model\n");
            return 0;
        }

        const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float);
        const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float);
        if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(low) < low_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal Q4 attention output batch received undersized buffers\n");
            return 0;
        }
        (void)group_tmp;
        (void)low_tmp;

        const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t);
        id<MTLBuffer> group_ids_buffer = nil;
        if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) {
            group_ids_buffer =
                ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids");
        } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer,
                                                 &g_attn_out_group_ids_bytes,
                                                 ids_bytes,
                                                 "ds4_attention_output_group_ids")) {
            group_ids_buffer = g_attn_out_group_ids_buffer;
        }
        if (!group_ids_buffer) return 0;

        int32_t *ids = (int32_t *)[group_ids_buffer contents];
        for (uint32_t t = 0; t < n_tokens; t++) {
            for (uint32_t group = 0; group < n_groups; group++) {
                ids[(uint64_t)t * n_groups + group] = (int32_t)group;
            }
        }

        uint64_t out_a_inner = 0;
        id<MTLBuffer> out_a_buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                     out_a_offset, out_a_bytes,
                                     &out_a_inner);
        if (!out_a_buf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        bool ok = true;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) ok = false;

        if (ok) {
            ds4_gpu_mul_mm_id_map_args map_args =
                ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim,
                                                n_groups,
                                                n_groups,
                                                n_groups,
                                                n_tokens);
            ds4_gpu_mul_mm_id_args mm_args =
                ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim,
                                            (uint32_t)rank,
                                            n_groups,
                                            row_a_bytes,
                                            (uint64_t)rank * row_a_bytes,
                                            n_groups,
                                            n_groups,
                                            n_tokens);
            id<MTLComputePipelineState> map_pipeline =
                ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups));
            id<MTLComputePipelineState> mm_pipeline =
                ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false);
            ok = ds4_gpu_encode_mul_mm_id(cb,
                                          map_pipeline,
                                          mm_pipeline,
                                          &map_args,
                                          &mm_args,
                                          out_a_buf,
                                          (NSUInteger)out_a_inner,
                                          ds4_gpu_tensor_buffer(heads),
                                          ds4_gpu_tensor_offset(heads),
                                          ds4_gpu_tensor_buffer(low),
                                          ds4_gpu_tensor_offset(low),
                                          group_ids_buffer,
                                          0) != 0;
        }

        if (ok) {
            ok = ds4_gpu_matmul_quant_tensor(out,
                                             model_map,
                                             model_size,
                                             out_b_offset,
                                             out_b_type,
                                             low_dim,
                                             out_dim,
                                             low,
                                             n_tokens) != 0;
        }

        if (!had_batch) {
            ok = ds4_gpu_end_commands() != 0 && ok;
        }
        return ok ? 1 : 0;
    }
}

int ds4_gpu_attention_output_q8_batch_f16_tensor(
        ds4_gpu_tensor       *out_h,
        ds4_gpu_tensor       *low,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                out_a_offset,
        uint64_t                out_b_offset,
        uint64_t                group_dim,
        uint64_t                rank,
        uint32_t                n_groups,
        uint64_t                out_dim,
        const ds4_gpu_tensor *heads,
        uint32_t                n_tokens) {
    (void)out_h; (void)low; (void)model_map; (void)model_size;
    (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank;
    (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens;
    return 0;
}

int ds4_gpu_matmul_q8_0_kslice_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                full_in_dim,
        uint64_t                k_off,
        uint64_t                k_cnt,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                x_elem_off) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if ((full_in_dim & 31u) != 0 || (k_off & 31u) != 0 || (k_cnt & 31u) != 0 ||
        k_cnt == 0 || k_off + k_cnt > full_in_dim ||
        full_in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) ||
            ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) {
            fprintf(stderr, "ds4: Metal Q8_0 kslice matmul received undersized buffers\n");
            return 0;
        }
        const uint64_t row_bytes = (full_in_dim / 32u) * 34u;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal Q8_0 kslice weights are outside the mapped model\n");
            return 0;
        }
        uint64_t inner = 0;
        id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                        weight_offset, weight_bytes, &inner);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        /* Same matvec kernel as the full projection: ne00 bounds the k loop
         * while nb01/nb02 keep the full-row stride, so each row reads only
         * the owned k window. */
        ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(full_in_dim, out_dim);
        mv_args.ne00 = (int32_t)k_cnt;
        mv_args.ne10 = (int32_t)k_cnt;
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        if (out_dim > 65536u) mv_dispatch.nsg = 8;
        mv_args.nr0 = mv_dispatch.nr0;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
        if (!pipeline) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
        [enc setBuffer:wbuf offset:(NSUInteger)(inner + (k_off / 32u) * 34u) atIndex:1];
        [enc setBuffer:xbuf
                offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float))
               atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
                                              1,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 kslice matvec")) {
            return 0;
        }
        return 1;
    }
}

int ds4_gpu_matmul_quant_kslice_tensor(
        ds4_gpu_tensor       *out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint32_t                weight_type,
        uint64_t                full_in_dim,
        uint64_t                k_off,
        uint64_t                k_cnt,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        uint64_t                x_elem_off) {
    if (weight_type == DS4_METAL_TENSOR_Q8_0) {
        return ds4_gpu_matmul_q8_0_kslice_tensor(out,
                                                 model_map,
                                                 model_size,
                                                 weight_offset,
                                                 full_in_dim,
                                                 k_off,
                                                 k_cnt,
                                                 out_dim,
                                                 x,
                                                 x_elem_off);
    }
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !x || !model_map ||
        full_in_dim == 0 || k_cnt == 0 || out_dim == 0 ||
        k_off + k_cnt > full_in_dim ||
        full_in_dim > UINT32_MAX || k_cnt > UINT32_MAX ||
        out_dim > UINT32_MAX) {
        return 0;
    }

    uint64_t block_elems = 0;
    uint64_t block_bytes = 0;
    if (weight_type == DS4_METAL_TENSOR_Q4_K) {
        block_elems = 256u;
        block_bytes = 144u;
    } else if (weight_type == DS4_METAL_TENSOR_Q4_0) {
        block_elems = 32u;
        block_bytes = 18u;
    } else {
        fprintf(stderr, "ds4: Metal quant kslice received unsupported type %u\n", weight_type);
        return 0;
    }
    if ((full_in_dim % block_elems) != 0 ||
        (k_off % block_elems) != 0 ||
        (k_cnt % block_elems) != 0) {
        fprintf(stderr, "ds4: Metal quant kslice dimensions are not block aligned\n");
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        if (!xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) ||
            ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) {
            fprintf(stderr, "ds4: Metal quant kslice matmul received undersized buffers\n");
            return 0;
        }

        const uint64_t row_bytes = (full_in_dim / block_elems) * block_bytes;
        const uint64_t weight_bytes = out_dim * row_bytes;
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal quant kslice weights are outside the mapped model\n");
            return 0;
        }
        uint64_t inner = 0;
        id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                        weight_offset, weight_bytes, &inner);
        if (!wbuf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (weight_type == DS4_METAL_TENSOR_Q4_K &&
            (k_cnt % 256u) == 0 &&
            getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) {
            const int16_t nsg = 2;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8);
            if (pipeline) {
                ds4_gpu_q8_0_matvec_args args = {
                    .ne00 = (int32_t)k_cnt,
                    .ne01 = (int32_t)out_dim,
                    .ne02 = 1,
                    .nb00 = 1,
                    .nb01 = row_bytes,
                    .nb02 = row_bytes * out_dim,
                    .nb03 = row_bytes * out_dim,
                    .ne10 = (int32_t)k_cnt,
                    .ne11 = 1,
                    .ne12 = 1,
                    .nb10 = sizeof(float),
                    .nb11 = k_cnt * sizeof(float),
                    .nb12 = k_cnt * sizeof(float),
                    .nb13 = k_cnt * sizeof(float),
                    .ne0 = (int32_t)out_dim,
                    .ne1 = 1,
                    .nr0 = 2,
                    .r2 = 1,
                    .r3 = 1,
                };
                const uint64_t rows_ptg = (uint64_t)nsg * 2u;
                const uint64_t w_skip = (k_off / block_elems) * block_bytes;

                id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBuffer:wbuf offset:(NSUInteger)(inner + w_skip) atIndex:1];
                [enc setBuffer:xbuf
                        offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float))
                       atIndex:2];
                [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
                [enc setThreadgroupMemoryLength:32 atIndex:0];
                [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg,
                                                      1,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);

                if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K kslice matvec")) return 0;
                return 1;
            }
        }

        fprintf(stderr, "ds4: Metal quant kslice has no kernel for type %u\n",
                weight_type);
        return 0;
    }
}

int ds4_gpu_attention_output_q8_tp_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *low,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                out_a_offset,
        uint64_t                out_b_offset,
        uint64_t                group_dim,
        uint64_t                rank,
        uint32_t                n_groups_total,
        uint32_t                group0,
        uint32_t                group_cnt,
        uint64_t                out_dim,
        const ds4_gpu_tensor *heads) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !low || !heads || !model_map ||
        group_dim == 0 || rank == 0 || group_cnt == 0 ||
        group0 + group_cnt > n_groups_total ||
        (group_dim % 32u) != 0 || ((rank * group_cnt) % 32u) != 0 ||
        group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t low_dim_total = (uint64_t)n_groups_total * rank;
        const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
        const uint64_t a_group_bytes = rank * row_a_bytes;
        const uint64_t out_a_bytes = (uint64_t)n_groups_total * a_group_bytes;
        if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) {
            fprintf(stderr, "ds4: Metal TP attention output weights are outside the mapped model\n");
            return 0;
        }
        /* The heads buffer holds only the owned groups, compact at its
         * base (the head slice keeps q/attention output halves packed). */
        if (ds4_gpu_tensor_bytes(heads) < (uint64_t)group_cnt * group_dim * sizeof(float) ||
            ds4_gpu_tensor_bytes(low) < (uint64_t)group_cnt * rank * sizeof(float) ||
            ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) {
            fprintf(stderr, "ds4: Metal TP attention output received undersized buffers\n");
            return 0;
        }

        uint64_t out_a_inner = 0;
        id<MTLBuffer> out_a_buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                       out_a_offset, out_a_bytes, &out_a_inner);
        if (!out_a_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) return 0;

        /* Low projection for the owned groups only: identical dispatch to
         * the single-node direct path with the weight base, heads input and
         * group count shifted to the slice.  The owned low half lands
         * compactly at low[0 .. group_cnt*rank). */
        ds4_gpu_mul_mv_id_args args = {
            .nei0 = (int32_t)group_cnt,
            .nei1 = 1,
            .nbi1 = 0,
            .ne00 = (int32_t)group_dim,
            .ne01 = (int32_t)rank,
            .ne02 = (int32_t)group_cnt,
            .nb00 = 34,
            .nb01 = row_a_bytes,
            .nb02 = a_group_bytes,
            .ne10 = (int32_t)group_dim,
            .ne11 = (int32_t)group_cnt,
            .ne12 = 1,
            .ne13 = 1,
            .nb10 = sizeof(float),
            .nb11 = (uint64_t)group_dim * sizeof(float),
            .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float),
            .ne0 = (int32_t)rank,
            .ne1 = (int32_t)group_cnt,
            .nb1 = (uint64_t)rank * sizeof(float),
            .nr0 = 2,
        };
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4);
        int ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
                pipeline,
                &args,
                out_a_buf,
                (NSUInteger)(out_a_inner + (uint64_t)group0 * a_group_bytes),
                ds4_gpu_tensor_buffer(heads),
                ds4_gpu_tensor_offset(heads),
                ds4_gpu_tensor_buffer(low),
                ds4_gpu_tensor_offset(low),
                32u * 2u * sizeof(float),
                4,
                true);
        if (!ok) return 0;

        /* Expand projection over the owned k window only; the result is this
         * rank's partial attention block output. */
        return ds4_gpu_matmul_q8_0_kslice_tensor(out, model_map, model_size,
                                                   out_b_offset,
                                                   low_dim_total,
                                                   (uint64_t)group0 * rank,
                                                   (uint64_t)group_cnt * rank,
                                                   out_dim, low, 0);
    }
}

int ds4_gpu_attention_output_low_q8_tensor(
        ds4_gpu_tensor       *low,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                out_a_offset,
        uint64_t                group_dim,
        uint64_t                rank,
        uint32_t                n_groups,
        const ds4_gpu_tensor *heads) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!low || !heads || !model_map || group_dim == 0 || rank == 0 ||
        n_groups == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t low_dim = (uint64_t)n_groups * rank;
        if ((group_dim % 32u) != 0 || low_dim > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal attention output low received invalid q8 dimensions\n");
            return 0;
        }

        const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
        const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
        if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) {
            fprintf(stderr, "ds4: Metal attention output low weights are outside the mapped model\n");
            return 0;
        }

        const uint64_t heads_bytes = (uint64_t)n_groups * group_dim * sizeof(float);
        const uint64_t low_bytes = low_dim * sizeof(float);
        if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(low) < low_bytes) {
            fprintf(stderr, "ds4: Metal attention output low received undersized buffers\n");
            return 0;
        }

        uint64_t out_a_inner = 0;
        id<MTLBuffer> out_a_buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                       out_a_offset, out_a_bytes,
                                       &out_a_inner);
        if (!out_a_buf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        bool ok = true;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) {
            ok = false;
        }

        if (ok) {
            ds4_gpu_mul_mv_id_args args = {
                .nei0 = (int32_t)n_groups,
                .nei1 = 1,
                .nbi1 = 0,
                .ne00 = (int32_t)group_dim,
                .ne01 = (int32_t)rank,
                .ne02 = (int32_t)n_groups,
                .nb00 = 34,
                .nb01 = row_a_bytes,
                .nb02 = (uint64_t)rank * row_a_bytes,
                .ne10 = (int32_t)group_dim,
                .ne11 = (int32_t)n_groups,
                .ne12 = 1,
                .ne13 = 1,
                .nb10 = sizeof(float),
                .nb11 = (uint64_t)group_dim * sizeof(float),
                .nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
                .ne0 = (int32_t)rank,
                .ne1 = (int32_t)n_groups,
                .nb1 = (uint64_t)rank * sizeof(float),
                .nr0 = 2,
            };
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4);
            ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
                                                         pipeline,
                                                         &args,
                                                         out_a_buf,
                                                         (NSUInteger)out_a_inner,
                                                         ds4_gpu_tensor_buffer(heads),
                                                         ds4_gpu_tensor_offset(heads),
                                                         ds4_gpu_tensor_buffer(low),
                                                         ds4_gpu_tensor_offset(low),
                                                         32u * 2u * sizeof(float),
                                                         4,
                                                         true) != 0;
        }

        if (!had_batch) {
            ok = ds4_gpu_end_commands() != 0 && ok;
        }
        return ok ? 1 : 0;
    }
}

int ds4_gpu_attention_output_low_q4_K_slice_tensor(
        ds4_gpu_tensor       *low,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                out_a_offset,
        uint64_t                group_dim,
        uint64_t                rank,
        uint32_t                group0,
        uint32_t                group_cnt,
        const ds4_gpu_tensor *heads) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!low || !heads || !model_map || group_dim == 0 || rank == 0 ||
        group_cnt == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        if ((group_dim % 256u) != 0) {
            fprintf(stderr, "ds4: Metal attention output low received invalid Q4_K dimensions\n");
            return 0;
        }

        uint64_t row_a_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K,
                                     (uint32_t)group_dim,
                                     &row_a_bytes)) {
            return 0;
        }
        if (rank > UINT64_MAX / row_a_bytes) return 0;
        const uint64_t group_weight_bytes = rank * row_a_bytes;
        if (group0 > UINT64_MAX / group_weight_bytes ||
            group_cnt > UINT64_MAX / group_weight_bytes) {
            return 0;
        }
        const uint64_t group_skip = (uint64_t)group0 * group_weight_bytes;
        const uint64_t out_a_bytes = (uint64_t)group_cnt * group_weight_bytes;
        if (out_a_offset > UINT64_MAX - group_skip ||
            out_a_offset + group_skip > model_size ||
            out_a_bytes > model_size - (out_a_offset + group_skip)) {
            fprintf(stderr, "ds4: Metal Q4 attention output low weights are outside the mapped model\n");
            return 0;
        }

        const uint64_t heads_bytes = (uint64_t)group_cnt * group_dim * sizeof(float);
        const uint64_t low_bytes = (uint64_t)group_cnt * rank * sizeof(float);
        if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(low) < low_bytes) {
            fprintf(stderr, "ds4: Metal Q4 attention output low received undersized buffers\n");
            return 0;
        }

        uint64_t out_a_inner = 0;
        id<MTLBuffer> out_a_buf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                      out_a_offset + group_skip,
                                      out_a_bytes,
                                      &out_a_inner);
        if (!out_a_buf) return 0;

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;

        bool ok = true;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb || owned) {
            ok = false;
        }

        if (ok) {
            ds4_gpu_mul_mv_id_args args = {
                .nei0 = (int32_t)group_cnt,
                .nei1 = 1,
                .nbi1 = 0,
                .ne00 = (int32_t)group_dim,
                .ne01 = (int32_t)rank,
                .ne02 = (int32_t)group_cnt,
                .nb00 = 1,
                .nb01 = row_a_bytes,
                .nb02 = group_weight_bytes,
                .ne10 = (int32_t)group_dim,
                .ne11 = (int32_t)group_cnt,
                .ne12 = 1,
                .ne13 = 1,
                .nb10 = sizeof(float),
                .nb11 = (uint64_t)group_dim * sizeof(float),
                .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float),
                .ne0 = (int32_t)rank,
                .ne1 = (int32_t)group_cnt,
                .nb1 = (uint64_t)rank * sizeof(float),
                .nr0 = 2,
            };
            const NSUInteger nsg = 2;
            id<MTLComputePipelineState> pipeline =
                ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q4_K_f32", (int16_t)nsg);
            ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
                                                         pipeline,
                                                         &args,
                                                         out_a_buf,
                                                         (NSUInteger)out_a_inner,
                                                         ds4_gpu_tensor_buffer(heads),
                                                         ds4_gpu_tensor_offset(heads),
                                                         ds4_gpu_tensor_buffer(low),
                                                         ds4_gpu_tensor_offset(low),
                                                         32u,
                                                         nsg,
                                                         false) != 0;
        }

        if (!had_batch) {
            ok = ds4_gpu_end_commands() != 0 && ok;
        }
        return ok ? 1 : 0;
    }
}

static NSUInteger ds4_gpu_align_up_ns(NSUInteger value, NSUInteger align) {
    return (value + align - 1u) & ~(align - 1u);
}

static int ds4_gpu_encode_cpy_f32_f32_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n) {
    if (!cb || !src || !dst || n == 0) return 0;

    ds4_gpu_cpy_args args =
        ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(float));
    const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f32_f32_pipeline);
    const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f32_f32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f32_f32_3d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint32_t             planes,
        uint64_t             src_row_stride,
        uint64_t             src_plane_stride,
        uint64_t             dst_row_stride,
        uint64_t             dst_plane_stride) {
    if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;

    ds4_gpu_cpy_args args = {
        .nk0 = (int64_t)cols,
        .ne00 = (int64_t)cols,
        .ne01 = (int64_t)rows,
        .ne02 = (int64_t)planes,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = src_row_stride,
        .nb02 = src_plane_stride,
        .nb03 = (uint64_t)planes * src_plane_stride,
        .ne0 = (int64_t)cols,
        .ne1 = (int64_t)rows,
        .ne2 = (int64_t)planes,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = dst_row_stride,
        .nb2 = dst_plane_stride,
        .nb3 = (uint64_t)planes * dst_plane_stride,
    };
    const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline);
    const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f32_f32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint32_t             planes,
        uint64_t             src_col_stride,
        uint64_t             src_row_stride,
        uint64_t             src_plane_stride,
        uint64_t             dst_row_stride,
        uint64_t             dst_plane_stride) {
    if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;

    ds4_gpu_cpy_args args = {
        .nk0 = (int64_t)cols,
        .ne00 = (int64_t)cols,
        .ne01 = (int64_t)rows,
        .ne02 = (int64_t)planes,
        .ne03 = 1,
        .nb00 = src_col_stride,
        .nb01 = src_row_stride,
        .nb02 = src_plane_stride,
        .nb03 = (uint64_t)planes * src_plane_stride,
        .ne0 = (int64_t)cols,
        .ne1 = (int64_t)rows,
        .ne2 = (int64_t)planes,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = dst_row_stride,
        .nb2 = dst_plane_stride,
        .nb3 = (uint64_t)planes * dst_plane_stride,
    };
    const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline);
    const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f32_f32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f32_f16_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n) {
    if (!cb || !src || !dst || n == 0) return 0;

    const int use_contiguous =
        ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0;
    id<MTLComputePipelineState> pipeline = use_contiguous
        ? g_cpy_contig_f32_f16_pipeline
        : g_cpy_f32_f16_pipeline;
    const NSUInteger work_items = use_contiguous
        ? ((NSUInteger)n + 3u) / 4u
        : (NSUInteger)n;
    const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline);
    const NSUInteger groups = (work_items + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    if (use_contiguous) {
        [enc setBytes:&n length:sizeof(n) atIndex:0];
    } else {
        ds4_gpu_cpy_args args =
            ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t));
        [enc setBytes:&args length:sizeof(args) atIndex:0];
    }
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f32_f16_2d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint64_t             src_row_stride,
        uint64_t             dst_row_stride) {
    if (!cb || !src || !dst || cols == 0 || rows == 0) return 0;

    ds4_gpu_cpy_args args = {
        .nk0 = (int64_t)cols,
        .ne00 = (int64_t)cols,
        .ne01 = (int64_t)rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = src_row_stride,
        .nb02 = (uint64_t)rows * src_row_stride,
        .nb03 = (uint64_t)rows * src_row_stride,
        .ne0 = (int64_t)cols,
        .ne1 = (int64_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(uint16_t),
        .nb1 = dst_row_stride,
        .nb2 = (uint64_t)rows * dst_row_stride,
        .nb3 = (uint64_t)rows * dst_row_stride,
    };
    const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline);
    const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f32_f16_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f32_f16_3d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint32_t             planes,
        uint64_t             src_row_stride,
        uint64_t             src_plane_stride,
        uint64_t             dst_row_stride,
        uint64_t             dst_plane_stride) {
    if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;

    ds4_gpu_cpy_args args = {
        .nk0 = (int64_t)cols,
        .ne00 = (int64_t)cols,
        .ne01 = (int64_t)rows,
        .ne02 = (int64_t)planes,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = src_row_stride,
        .nb02 = src_plane_stride,
        .nb03 = (uint64_t)planes * src_plane_stride,
        .ne0 = (int64_t)cols,
        .ne1 = (int64_t)rows,
        .ne2 = (int64_t)planes,
        .ne3 = 1,
        .nb0 = sizeof(uint16_t),
        .nb1 = dst_row_stride,
        .nb2 = dst_plane_stride,
        .nb3 = (uint64_t)planes * dst_plane_stride,
    };
    const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline);
    const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f32_f16_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f16_f16_3d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             cols,
        uint32_t             rows,
        uint32_t             planes,
        uint64_t             src_row_stride,
        uint64_t             src_plane_stride,
        uint64_t             dst_row_stride,
        uint64_t             dst_plane_stride) {
    if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;

    ds4_gpu_cpy_args args = {
        .nk0 = (int64_t)cols,
        .ne00 = (int64_t)cols,
        .ne01 = (int64_t)rows,
        .ne02 = (int64_t)planes,
        .ne03 = 1,
        .nb00 = sizeof(uint16_t),
        .nb01 = src_row_stride,
        .nb02 = src_plane_stride,
        .nb03 = (uint64_t)planes * src_plane_stride,
        .ne0 = (int64_t)cols,
        .ne1 = (int64_t)rows,
        .ne2 = (int64_t)planes,
        .ne3 = 1,
        .nb0 = sizeof(uint16_t),
        .nb1 = dst_row_stride,
        .nb2 = dst_plane_stride,
        .nb3 = (uint64_t)planes * dst_plane_stride,
    };
    const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f16_f16_pipeline);
    const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f16_f16_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_cpy_f16_f32_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n) {
    if (!cb || !src || !dst || n == 0) return 0;

    const int use_contiguous =
        ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0;
    id<MTLComputePipelineState> pipeline = use_contiguous
        ? g_cpy_contig_f16_f32_pipeline
        : g_cpy_f16_f32_pipeline;
    const NSUInteger work_items = use_contiguous
        ? ((NSUInteger)n + 3u) / 4u
        : (NSUInteger)n;
    const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline);
    const NSUInteger groups = (work_items + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    if (use_contiguous) {
        [enc setBytes:&n length:sizeof(n) atIndex:0];
    } else {
        ds4_gpu_cpy_args args =
            ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float));
        [enc setBytes:&args length:sizeof(args) atIndex:0];
    }
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_copy_to_f16_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        bool                 src_is_f16,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             n) {
    if (!cb || !src || !dst) return 0;
    if (n == 0) return 1;
    if (!src_is_f16) {
        return ds4_gpu_encode_cpy_f32_f16_1d(cb, src, src_off, dst, dst_off, n);
    }

    if (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") <= 0) {
        const NSUInteger work_items = ((NSUInteger)n + 3u) / 4u;
        const NSUInteger nth = ds4_gpu_cpy_threads(
                (uint32_t)work_items,
                g_cpy_contig_f16_f16_pipeline);
        const NSUInteger groups = (work_items + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_cpy_contig_f16_f16_pipeline];
        [enc setBytes:&n length:sizeof(n) atIndex:0];
        [enc setBuffer:src offset:src_off atIndex:1];
        [enc setBuffer:dst offset:dst_off atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        return 1;
    }

    if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder();
    id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
    if (!blit) return 0;
    [blit copyFromBuffer:src
            sourceOffset:src_off
                toBuffer:dst
       destinationOffset:dst_off
                    size:(NSUInteger)n * sizeof(uint16_t)];
    [blit endEncoding];
    return 1;
}

static int ds4_gpu_encode_copy_raw_ring_to_f16(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        raw,
        NSUInteger           raw_offset,
        uint32_t             raw_cap,
        uint32_t             raw_start,
        uint32_t             n_raw,
        uint32_t             head_dim,
        id<MTLBuffer>        dst,
        NSUInteger           dst_offset) {
    if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap ||
        n_raw == 0 || n_raw > raw_cap || head_dim == 0) {
        return 0;
    }

    const NSUInteger raw_row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger dst_row_bytes = (NSUInteger)head_dim * sizeof(uint16_t);
    const uint32_t tail_rows = raw_cap - raw_start < n_raw
        ? raw_cap - raw_start
        : n_raw;
    const uint32_t head_rows = n_raw - tail_rows;
    const uint64_t tail_count = (uint64_t)tail_rows * head_dim;
    const uint64_t head_count = (uint64_t)head_rows * head_dim;
    const uint64_t raw_inner = (uint64_t)raw_start * raw_row_bytes;
    const uint64_t dst_inner = (uint64_t)tail_rows * dst_row_bytes;
    if (tail_count > UINT32_MAX || head_count > UINT32_MAX ||
        raw_inner > NSUIntegerMax - raw_offset ||
        dst_inner > NSUIntegerMax - dst_offset) {
        return 0;
    }

    if (tail_rows &&
        !ds4_gpu_encode_cpy_f32_f16_1d(
                cb,
                raw,
                raw_offset + (NSUInteger)raw_inner,
                dst,
                dst_offset,
                (uint32_t)tail_count)) {
        return 0;
    }
    if (head_rows &&
        !ds4_gpu_encode_cpy_f32_f16_1d(
                cb,
                raw,
                raw_offset,
                dst,
                dst_offset + (NSUInteger)dst_inner,
                (uint32_t)head_count)) {
        return 0;
    }
    return 1;
}

static int ds4_gpu_encode_flash_kv_stage_f16(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        raw,
        NSUInteger           raw_offset,
        uint32_t             raw_cap,
        uint32_t             raw_start,
        uint32_t             n_raw,
        id<MTLBuffer>        comp,
        NSUInteger           comp_offset,
        bool                 comp_is_f16,
        uint32_t             n_comp,
        uint32_t             head_dim,
        id<MTLBuffer>        dst,
        NSUInteger           dst_offset,
        id<MTLBuffer>        mask,
        NSUInteger           mask_offset,
        id<MTLBuffer>        pad,
        NSUInteger           pad_offset,
        bool                 fuse_pad,
        bool                 shared_pad,
        bool                *did_fuse_pad) {
    if (did_fuse_pad) *did_fuse_pad = false;
    if (!cb || !raw || !comp || !dst || raw_cap == 0 ||
        raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap ||
        n_comp == 0 || head_dim == 0) {
        return 0;
    }

    const bool force =
        getenv("DS4_METAL_ENABLE_GATHERED_KV_STAGE") != NULL;
    const bool disabled =
        getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE") != NULL;
    const bool require =
        getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL;
    const bool supported_shape =
        comp_is_f16 && head_dim == 512u && raw_cap <= UINT32_MAX / 128u;
    const uint64_t row_vecs64 = 128u;
    const uint64_t comp_count64 = (uint64_t)n_comp * head_dim;
    const uint64_t dst_comp_inner64 =
        (uint64_t)n_raw * head_dim * sizeof(uint16_t);
    if (comp_count64 > UINT32_MAX ||
        dst_comp_inner64 > NSUIntegerMax - dst_offset) {
        return 0;
    }
    const uint64_t total_vecs64 =
        ((uint64_t)n_raw + n_comp) * row_vecs64;
    const bool valid_grid =
        total_vecs64 != 0 && total_vecs64 <= UINT32_MAX;
    const bool eligible =
        supported_shape && valid_grid && !g_quality_mode && !disabled &&
        g_flash_kv_stage_f16_pipeline != nil &&
        (ds4_gpu_device_name_contains("M3") ||
         ds4_gpu_device_name_contains("M5") || force);
    const bool component_disabled = eligible &&
        (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 ||
         ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0);
    const bool use_fusion = eligible && !component_disabled;
    const bool use_pad_fusion =
        use_fusion && fuse_pad && mask != nil && pad != nil &&
        getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION") == NULL;
    if (require && supported_shape && !use_fusion) {
        fprintf(stderr,
                "ds4: required Metal gathered KV staging kernel was not selected\n");
        return 0;
    }

    if (use_fusion) {
        ds4_gpu_flash_kv_stage_f16_args args = {
            .raw_cap = raw_cap,
            .raw_start = raw_start,
            .n_raw = n_raw,
            .n_comp = n_comp,
            .pad_rows = use_pad_fusion ? 32u : 0u,
            .shared_pad = use_pad_fusion && shared_pad ? 1u : 0u,
        };
        const NSUInteger total_vecs = (NSUInteger)total_vecs64 +
            (use_pad_fusion ? 32u * 128u + 32u : 0u);
        NSUInteger nth =
            g_flash_kv_stage_f16_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > 256u) nth = 256u;
        if (nth > total_vecs) nth = total_vecs;
        if (nth == 0) return 0;
        const NSUInteger groups = (total_vecs + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_flash_kv_stage_f16_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:raw offset:raw_offset atIndex:1];
        [enc setBuffer:comp offset:comp_offset atIndex:2];
        [enc setBuffer:dst offset:dst_offset atIndex:3];
        [enc setBuffer:(use_pad_fusion ? mask : dst)
                 offset:(use_pad_fusion ? mask_offset : dst_offset)
                atIndex:4];
        [enc setBuffer:(use_pad_fusion ? pad : dst)
                 offset:(use_pad_fusion ? pad_offset : dst_offset)
                atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        if (did_fuse_pad) *did_fuse_pad = use_pad_fusion;
        return 1;
    }

    if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
                                              raw,
                                              raw_offset,
                                              raw_cap,
                                              raw_start,
                                              n_raw,
                                              head_dim,
                                              dst,
                                              dst_offset)) {
        return 0;
    }
    return ds4_gpu_encode_copy_to_f16_1d(
        cb,
        comp,
        comp_offset,
        comp_is_f16,
        dst,
        dst_offset + (NSUInteger)dst_comp_inner64,
        (uint32_t)comp_count64);
}

int ds4_gpu_flash_kv_stage_f16_tensor(
        ds4_gpu_tensor       *dst,
        const ds4_gpu_tensor *raw,
        uint32_t                raw_cap,
        uint32_t                raw_start,
        uint32_t                n_raw,
        const ds4_gpu_tensor *comp,
        uint32_t                comp_is_f16,
        uint32_t                n_comp,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!dst || !raw || !comp || raw_cap == 0 || raw_start >= raw_cap ||
        n_raw == 0 || n_raw > raw_cap || n_comp == 0 ||
        comp_is_f16 == 0 || head_dim != 512u) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t raw_bytes =
            (uint64_t)raw_cap * head_dim * sizeof(float);
        const uint64_t comp_bytes =
            (uint64_t)n_comp * head_dim *
            (comp_is_f16 ? sizeof(uint16_t) : sizeof(float));
        const uint64_t dst_bytes =
            ((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp);
        id<MTLBuffer> dstbuf = ds4_gpu_tensor_buffer(dst);
        if (!rawbuf || !compbuf || !dstbuf ||
            ds4_gpu_tensor_bytes(raw) < raw_bytes ||
            ds4_gpu_tensor_bytes(comp) < comp_bytes ||
            ds4_gpu_tensor_bytes(dst) < dst_bytes) {
            fprintf(stderr,
                    "ds4: Metal gathered KV staging received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;
        if (!ds4_gpu_encode_flash_kv_stage_f16(
                cb,
                rawbuf,
                ds4_gpu_tensor_offset(raw),
                raw_cap,
                raw_start,
                n_raw,
                compbuf,
                ds4_gpu_tensor_offset(comp),
                comp_is_f16 != 0,
                n_comp,
                head_dim,
                dstbuf,
                ds4_gpu_tensor_offset(dst),
                nil,
                0,
                nil,
                0,
                false,
                false,
                NULL)) {
            return 0;
        }
        if (!ds4_gpu_finish_command_buffer(
                cb, owned, "gathered KV staging")) {
            return 0;
        }
    }
    return 1;
}

static int ds4_gpu_encode_fill_f16_1d(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        buf,
        NSUInteger           offset,
        uint32_t             n,
        float                value) {
    if (!cb || !buf || n == 0) return 0;

    ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(n, 1, 0, 0.0f, 0.0f);
    args.val = value;

    NSUInteger nth = (NSUInteger)n;
    const NSUInteger max_threads = g_unary_fill_f16_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth > 256u) nth = 256u;
    if (nth == 0) nth = 1u;
    const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_unary_fill_f16_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:buf offset:offset atIndex:1];
    [enc setBuffer:buf offset:offset atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_flash_attention_raw_heads(
        id<MTLCommandBuffer>  cb,
        ds4_gpu_tensor     *heads,
        id<MTLBuffer>         sinks_buf,
        NSUInteger            sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t              n_raw,
        uint32_t              raw_cap,
        uint32_t              raw_start,
        uint32_t              n_head,
        uint32_t              head_dim) {
    if (head_dim != 512 || n_head == 0 || n_raw == 0 || raw_cap < n_raw) {
        return 0;
    }

    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
    const uint64_t heads_bytes = q_bytes;
    if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        ds4_gpu_tensor_bytes(heads) < heads_bytes) {
        fprintf(stderr, "ds4: Metal DS4 FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t ncpsg = 32;
    const uint32_t nwg = 32;
    const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_raw, nwg, ncpsg);
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_raw * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16;
    const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 +
                                 (NSUInteger)ncpsg * sizeof(uint16_t);
    const NSUInteger nrows = (NSUInteger)n_head;
    const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
                                 nrows * (2u * (NSUInteger)nwg) * sizeof(float);

    id<MTLBuffer> mask_buffer =
        ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv_f16") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
                                         &g_flash_attn_tmp_bytes,
                                         tmp_bytes,
                                         "ds4_flash_attn_tmp")) {
        return 0;
    }
    memset([mask_buffer contents], 0, mask_bytes);

    id<MTLComputePipelineState> pad_pipeline = nil;
    if ((n_raw % ncpsg) != 0) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> vec_pipeline =
        ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
                                              true, true, false, false, (n_raw % ncpsg) != 0,
                                              false,
                                              (int32_t)head_dim,
                                              (int32_t)head_dim,
                                              (int32_t)nsg,
                                              (int32_t)nwg);
    id<MTLComputePipelineState> reduce_pipeline =
        ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
    if (!vec_pipeline || !reduce_pipeline) return 0;

    if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              raw_cap,
                                              raw_start,
                                              n_raw,
                                              head_dim,
                                              g_flash_attn_kv_buffer,
                                              0)) {
        return 0;
    }

    if ((n_raw % ncpsg) != 0) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_raw,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_raw * row_bytes_f16,
            .nb13 = (uint64_t)n_raw * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_raw * row_bytes_f16,
            .nb23 = (uint64_t)n_raw * row_bytes_f16,
            .ne31 = 1,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = mask_bytes,
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
    }

    ds4_gpu_flash_attn_vec_args vec_args = {
        .ne01 = 1,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_head * row_bytes,
        .ne11 = (int32_t)n_raw,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_raw * row_bytes_f16,
        .nb13 = (uint64_t)n_raw * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_raw * row_bytes_f16,
        .nb23 = (uint64_t)n_raw * row_bytes_f16,
        .ne31 = 1,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = mask_bytes,
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = 1,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
                                     4u * ncpsg +
                                     2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:vec_pipeline];
    [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    ds4_gpu_flash_attn_reduce_args reduce_args = {
        .nrows = (int32_t)nrows,
    };
    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:reduce_pipeline];
    [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

/* Rectangular causal/window mask for raw prefill: q covers rows
 * [q_row0, q_row0 + n_q) of an n_kv-token block whose keys all live in
 * rows [0, n_kv).  The square prefill case is q_row0 == 0, n_q == n_kv. */
static void ds4_gpu_fill_raw_prefill_mask(
        uint16_t *mask,
        uint32_t  q_row0,
        uint32_t  n_q,
        uint32_t  n_kv,
        uint32_t  window) {
    const uint16_t neg_inf_half = 0xfc00u;
    for (uint32_t q = 0; q < n_q; q++) {
        const uint32_t qpos = q_row0 + q;
        uint16_t *row = mask + (uint64_t)q * n_kv;
        for (uint32_t k = 0; k < n_kv; k++) {
            const bool causal = k <= qpos;
            const bool in_window = window == 0 || qpos - k < window;
            row[k] = causal && in_window ? 0u : neg_inf_half;
        }
    }
}

static void ds4_gpu_fill_glm_prefill_mask(
        uint16_t *mask,
        uint32_t  pos0,
        uint32_t  n_tokens,
        uint32_t  cache_len) {
    const uint16_t neg_inf_half = 0xfc00u;
    for (uint32_t q = 0; q < n_tokens; q++) {
        const uint32_t qpos = pos0 + q;
        uint16_t *row = mask + (uint64_t)q * cache_len;
        for (uint32_t k = 0; k < cache_len; k++) {
            row[k] = k <= qpos ? 0u : neg_inf_half;
        }
    }
}

static id<MTLBuffer> ds4_gpu_glm_prefill_mask_buffer(
        uint32_t pos0,
        uint32_t n_tokens,
        uint32_t cache_len,
        NSUInteger mask_bytes) {
    const int same_shape =
        g_glm_flash_attn_mask_valid &&
        g_glm_flash_attn_mask_buffer &&
        g_glm_flash_attn_mask_bytes >= mask_bytes &&
        g_glm_flash_attn_mask_pos0 == pos0 &&
        g_glm_flash_attn_mask_tokens == n_tokens &&
        g_glm_flash_attn_mask_cache_len == cache_len;
    if (same_shape) return g_glm_flash_attn_mask_buffer;

    if (g_glm_flash_attn_mask_buffer) {
        [g_transient_buffers addObject:g_glm_flash_attn_mask_buffer];
        g_glm_flash_attn_mask_buffer = nil;
    }
    g_glm_flash_attn_mask_bytes = 0;
    g_glm_flash_attn_mask_valid = 0;
    if (!ds4_gpu_ensure_scratch_buffer(&g_glm_flash_attn_mask_buffer,
                                        &g_glm_flash_attn_mask_bytes,
                                        mask_bytes,
                                        "ds4_glm_flash_attn_mask")) {
        return nil;
    }

    ds4_gpu_fill_glm_prefill_mask((uint16_t *)[g_glm_flash_attn_mask_buffer contents],
                                  pos0,
                                  n_tokens,
                                  cache_len);
    g_glm_flash_attn_mask_pos0 = pos0;
    g_glm_flash_attn_mask_tokens = n_tokens;
    g_glm_flash_attn_mask_cache_len = cache_len;
    g_glm_flash_attn_mask_valid = 1;
    return g_glm_flash_attn_mask_buffer;
}

static void ds4_gpu_fill_raw_decode_batch_mask(
        uint16_t *mask,
        uint32_t  n_tokens,
        uint32_t  n_raw,
        uint32_t  pos0,
        uint32_t  window) {
    const uint16_t neg_inf_half = 0xfc00u;
    const uint32_t last_pos = pos0 + n_tokens - 1u;
    /* The caller has already copied the SWA ring into logical order when it
     * wraps, so key row k represents first_raw_pos + k. */
    const uint32_t first_raw_pos = last_pos + 1u - n_raw;
    for (uint32_t q = 0; q < n_tokens; q++) {
        const uint32_t qpos = pos0 + q;
        uint16_t *row = mask + (uint64_t)q * n_raw;
        for (uint32_t k = 0; k < n_raw; k++) {
            const uint32_t kpos = first_raw_pos + k;
            const bool causal = kpos <= qpos;
            const bool in_window = causal && (window == 0 || qpos - kpos < window);
            row[k] = causal && in_window ? 0u : neg_inf_half;
        }
    }
}

static void ds4_gpu_fill_mixed_decode_batch_mask(
        uint16_t *mask,
        uint32_t  n_tokens,
        uint32_t  n_raw,
        uint32_t  n_comp,
        uint32_t  pos0,
        uint32_t  window,
        uint32_t  ratio) {
    const uint16_t neg_inf_half = 0xfc00u;
    const uint32_t n_keys = n_raw + n_comp;
    const uint32_t last_pos = pos0 + n_tokens - 1u;
    /* Raw keys are laid out by logical position; compressed keys follow them. */
    const uint32_t first_raw_pos = last_pos + 1u - n_raw;
    for (uint32_t q = 0; q < n_tokens; q++) {
        const uint32_t qpos = pos0 + q;
        uint16_t *row = mask + (uint64_t)q * n_keys;
        for (uint32_t k = 0; k < n_raw; k++) {
            const uint32_t kpos = first_raw_pos + k;
            const bool causal = kpos <= qpos;
            const bool in_window = causal && (window == 0 || qpos - kpos < window);
            row[k] = causal && in_window ? 0u : neg_inf_half;
        }
        const uint32_t n_visible = (qpos + 1u) / ratio;
        for (uint32_t c = 0; c < n_comp; c++) {
            row[n_raw + c] = c < n_visible ? 0u : neg_inf_half;
        }
    }
}

/* Rectangular causal/window + compressed-key visibility mask: q covers rows
 * [q_row0, q_row0 + n_q) of an n_tokens-token chunk whose raw keys all stay
 * resident, followed by n_comp compressed keys.  The square prefill case is
 * q_row0 == 0, n_q == n_tokens. */
static void ds4_gpu_fill_static_mixed_prefill_mask(
        uint16_t *mask,
        uint32_t  q_row0,
        uint32_t  n_q,
        uint32_t  n_tokens,
        uint32_t  n_comp,
        uint32_t  window,
        uint32_t  ratio) {
    const uint16_t neg_inf_half = 0xfc00u;
    const uint32_t n_keys = n_tokens + n_comp;
    for (uint32_t q = 0; q < n_q; q++) {
        const uint32_t qpos = q_row0 + q;
        uint16_t *row = mask + (uint64_t)q * n_keys;
        for (uint32_t k = 0; k < n_tokens; k++) {
            const bool causal = k <= qpos;
            const bool in_window = window == 0 || qpos - k < window;
            row[k] = causal && in_window ? 0u : neg_inf_half;
        }

        const uint32_t n_visible = (qpos + 1u) / ratio;
        for (uint32_t c = 0; c < n_comp; c++) {
            row[n_tokens + c] = c < n_visible ? 0u : neg_inf_half;
        }
    }
}

/* Static-mixed prefill FlashAttention over a rectangular problem: q holds
 * n_q query rows for token positions [q_row0, q_row0 + n_q) of the chunk,
 * while the keys stay full (all n_tokens raw rows plus n_comp compressed
 * rows).  The classic square prefill is q_row0 == 0, n_q == n_tokens. */
static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(
        id<MTLCommandBuffer> __strong *cbp,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t               comp_kv_f16,
        const ds4_gpu_tensor *comp_mask,
        uint32_t               use_comp_mask,
        uint32_t               q_row0,
        uint32_t               n_q,
        uint32_t               n_tokens,
        uint32_t               n_comp,
        uint32_t               window,
        uint32_t               ratio,
        uint32_t               n_head,
        uint32_t               head_dim) {
    if (!cbp || !*cbp) return 0;
    id<MTLCommandBuffer> cb = *cbp;
    if (head_dim != 512 || n_head == 0 || n_q == 0 || n_tokens == 0 || ratio == 0) {
        return 0;
    }

    const uint32_t n_keys = n_tokens + n_comp;
    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf;
    id<MTLBuffer> maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
    const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
                                (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
    const uint64_t comp_mask_bytes = use_comp_mask
        ? (uint64_t)n_comp * (q_row0 + n_q) * sizeof(float) : 0u;
    if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
        (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
        ds4_gpu_tensor_bytes(heads) < q_bytes) {
        fprintf(stderr, "ds4: Metal prefill static mixed DS4 non-vector FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t nqptg = 8;
    const uint32_t ncpsg = 64;
    const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
    const bool has_kvpad = (n_keys % ncpsg) != 0;
    const bool bc_mask = (n_q % nqptg) != 0;
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_q * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
    const NSUInteger pad_bytes = has_kvpad
        ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t))
        : 1u;
    const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg;
    const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg;
    const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);

    const uint32_t mask_cache_kind = ratio == 4u
        ? DS4_GPU_PREFILL_MASK_CACHE_RATIO4
        : (ratio == 128u ? DS4_GPU_PREFILL_MASK_CACHE_RATIO128 : 0u);
    bool mask_cache_created = false;
    ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache =
        use_comp_mask == 0u && mask_cache_kind != 0u &&
        q_row0 == 0u && n_q == n_tokens
            ? ds4_gpu_get_zero_prefix_prefill_mask_cache(mask_cache_kind,
                                                          n_tokens,
                                                          n_comp,
                                                          n_keys,
                                                          window,
                                                          ratio,
                                                          nqptg,
                                                          ncpsg,
                                                          has_kvpad,
                                                          bc_mask,
                                                          mask_bytes,
                                                          blk_bytes,
                                                          &mask_cache_created)
            : NULL;
    id<MTLBuffer> mask_buffer = mask_cache
        ? mask_cache->mask
        : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (mask_cache && mask_cache_created) {
        ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
                                                 0u,
                                                 n_tokens,
                                                 n_tokens,
                                                 n_comp,
                                                 window,
                                                 ratio);
        mask_cache->valid = true;
    }
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv_f16") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        (!mask_cache &&
         !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
                                          &g_flash_attn_blk_bytes,
                                          blk_bytes,
                                          "ds4_flash_attn_blk"))) {
        return 0;
    }
    id<MTLBuffer> blk_buffer = mask_cache
        ? mask_cache->blk
        : g_flash_attn_blk_buffer;

    const bool flash_stage_profile =
        getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
    double flash_stage_t0 = 0.0;
    if (flash_stage_profile) {
        if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
            return 0;
        }
        int profile_owned = 0;
        cb = ds4_gpu_command_buffer(&profile_owned);
        if (!cb || profile_owned) return 0;
        *cbp = cb;
        flash_stage_t0 = ds4_gpu_now_ms();
    }
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
        if (flash_stage_profile) { \
            if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
                    "static_mixed_nonvec", (name), n_q, n_comp, n_keys, \
                    n_head, head_dim, window, ratio, &flash_stage_t0)) { \
                return 0; \
            } \
            cb = *cbp; \
        } \
    } while (0)

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         rawbuf,
                                         ds4_gpu_tensor_offset(raw_kv),
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_tokens * head_dim)) {
        return 0;
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");
    if (n_comp &&
        !ds4_gpu_encode_copy_to_f16_1d(cb,
                                       compbuf,
                                       ds4_gpu_tensor_offset(comp_kv),
                                       comp_kv_f16 != 0,
                                       g_flash_attn_kv_buffer,
                                       (NSUInteger)n_tokens * row_bytes_f16,
                                       n_comp * head_dim)) {
        return 0;
    }
    if (n_comp) {
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp");
    }

    if (!mask_cache) {
        ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
                                                 q_row0,
                                                 n_q,
                                                 n_tokens,
                                                 n_comp,
                                                 window,
                                                 ratio);
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");
    if (use_comp_mask && n_comp != 0) {
        if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
                                             maskbuf,
                                             ds4_gpu_tensor_offset(comp_mask) +
                                                 (NSUInteger)((uint64_t)q_row0 * n_comp * sizeof(float)),
                                             mask_buffer,
                                             (NSUInteger)n_tokens * sizeof(uint16_t),
                                             n_comp,
                                             n_q,
                                             (uint64_t)n_comp * sizeof(float),
                                             (uint64_t)n_keys * sizeof(uint16_t))) {
            return 0;
        }
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy");
    }

    id<MTLComputePipelineState> pad_pipeline = nil;
    if (has_kvpad) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> blk_pipeline =
        ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
    id<MTLComputePipelineState> attn_pipeline =
        ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
                                          true, true, false, false, has_kvpad, bc_mask,
                                          (int32_t)head_dim,
                                          (int32_t)head_dim,
                                          (int32_t)nsg);
    if (!blk_pipeline || !attn_pipeline) return 0;

    if (has_kvpad) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_keys,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_keys * row_bytes_f16,
            .nb13 = (uint64_t)n_keys * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_keys * row_bytes_f16,
            .nb23 = (uint64_t)n_keys * row_bytes_f16,
            .ne31 = (int32_t)n_q,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
    }

    ds4_gpu_flash_attn_blk_args blk_args = {
        .ne01 = (int32_t)n_q,
        .ne30 = (int32_t)n_keys,
        .ne31 = (int32_t)n_q,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
    };

    id<MTLComputeCommandEncoder> enc = nil;
    if (!mask_cache || !mask_cache->blk_ready) {
        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:blk_pipeline];
        [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
        [enc setBuffer:mask_buffer offset:0 atIndex:1];
        [enc setBuffer:blk_buffer offset:0 atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        if (mask_cache) mask_cache->blk_ready = true;
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map");

    ds4_gpu_flash_attn_vec_args args = {
        .ne01 = (int32_t)n_q,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_q * n_head * row_bytes,
        .ne11 = (int32_t)n_keys,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_keys * row_bytes_f16,
        .nb13 = (uint64_t)n_keys * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_keys * row_bytes_f16,
        .nb23 = (uint64_t)n_keys * row_bytes_f16,
        .ne31 = (int32_t)n_q,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = (int32_t)n_q,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
    const NSUInteger shared_elems = (NSUInteger)nqptg *
        ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:attn_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:blk_buffer offset:0 atIndex:7];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention");

#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
    return 1;
}

static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(
        id<MTLCommandBuffer> __strong *cbp,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t               comp_kv_f16,
        const ds4_gpu_tensor *comp_mask,
        uint32_t               use_comp_mask,
        uint32_t               n_tokens,
        uint32_t               n_comp,
        uint32_t               window,
        uint32_t               ratio,
        uint32_t               n_head,
        uint32_t               head_dim) {
    if (!cbp || !*cbp) return 0;
    id<MTLCommandBuffer> cb = *cbp;
    if (head_dim != 512 || n_head == 0 || n_tokens == 0 || ratio == 0) {
        return 0;
    }

    const uint32_t n_keys = n_tokens + n_comp;
    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf;
    id<MTLBuffer> maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
    const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
                                (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
    const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u;
    if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
        (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
        ds4_gpu_tensor_bytes(heads) < q_bytes) {
        fprintf(stderr, "ds4: Metal prefill static mixed DS4 FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t ncpsg = 32;
    const uint32_t nwg = 32;
    const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg);
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
    const bool has_kvpad = (n_keys % ncpsg) != 0;
    const NSUInteger pad_bytes = has_kvpad
        ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t))
        : 1u;
    const NSUInteger nrows = (NSUInteger)n_tokens * n_head;
    const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
                                 nrows * (2u * (NSUInteger)nwg) * sizeof(float);

    id<MTLBuffer> mask_buffer =
        ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
                                         &g_flash_attn_tmp_bytes,
                                         tmp_bytes,
                                         "ds4_flash_attn_tmp")) {
        return 0;
    }

    const bool flash_stage_profile =
        getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
    double flash_stage_t0 = 0.0;
    if (flash_stage_profile) {
        if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
            return 0;
        }
        int profile_owned = 0;
        cb = ds4_gpu_command_buffer(&profile_owned);
        if (!cb || profile_owned) return 0;
        *cbp = cb;
        flash_stage_t0 = ds4_gpu_now_ms();
    }
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
        if (flash_stage_profile) { \
            if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
                    "static_mixed_vec", (name), n_tokens, n_comp, n_keys, \
                    n_head, head_dim, window, ratio, &flash_stage_t0)) { \
                return 0; \
            } \
            cb = *cbp; \
        } \
    } while (0)

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         rawbuf,
                                         ds4_gpu_tensor_offset(raw_kv),
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_tokens * head_dim)) {
        return 0;
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");
    if (n_comp) {
        if (!ds4_gpu_encode_copy_to_f16_1d(cb,
                                           compbuf,
                                           ds4_gpu_tensor_offset(comp_kv),
                                           comp_kv_f16 != 0,
                                           g_flash_attn_kv_buffer,
                                           (NSUInteger)n_tokens * row_bytes_f16,
                                           n_comp * head_dim)) {
            return 0;
        }
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp");
    }

    ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
                                             0,
                                             n_tokens,
                                             n_tokens,
                                             n_comp,
                                             window,
                                             ratio);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");
    if (use_comp_mask && n_comp != 0) {
        if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
                                             maskbuf,
                                             ds4_gpu_tensor_offset(comp_mask),
                                             mask_buffer,
                                             (NSUInteger)n_tokens * sizeof(uint16_t),
                                             n_comp,
                                             n_tokens,
                                             (uint64_t)n_comp * sizeof(float),
                                             (uint64_t)n_keys * sizeof(uint16_t))) {
            return 0;
        }
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy");
    }

    id<MTLComputePipelineState> pad_pipeline = nil;
    id<MTLComputeCommandEncoder> enc = nil;
    if (has_kvpad) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> vec_pipeline =
        ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
                                              true, true, false, false, has_kvpad,
                                              false,
                                              (int32_t)head_dim,
                                              (int32_t)head_dim,
                                              (int32_t)nsg,
                                              (int32_t)nwg);
    id<MTLComputePipelineState> reduce_pipeline =
        ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
    if (!vec_pipeline || !reduce_pipeline) return 0;

    if (has_kvpad) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_keys,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_keys * row_bytes_f16,
            .nb13 = (uint64_t)n_keys * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_keys * row_bytes_f16,
            .nb23 = (uint64_t)n_keys * row_bytes_f16,
            .ne31 = (int32_t)n_tokens,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
    }

    ds4_gpu_flash_attn_vec_args vec_args = {
        .ne01 = (int32_t)n_tokens,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_tokens * n_head * row_bytes,
        .ne11 = (int32_t)n_keys,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_keys * row_bytes_f16,
        .nb13 = (uint64_t)n_keys * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_keys * row_bytes_f16,
        .nb23 = (uint64_t)n_keys * row_bytes_f16,
        .ne31 = (int32_t)n_tokens,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = (int32_t)n_tokens,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
                                     4u * ncpsg +
                                     2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:vec_pipeline];
    [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec");

    ds4_gpu_flash_attn_reduce_args reduce_args = {
        .nrows = (int32_t)nrows,
    };
    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:reduce_pipeline];
    [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce");

#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
    return 1;
}

static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(
        id<MTLCommandBuffer> __strong *cbp,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t               comp_kv_f16,
        const ds4_gpu_tensor *comp_mask,
        uint32_t               use_comp_mask,
        uint32_t               q_row0,
        uint32_t               n_q,
        uint32_t               n_tokens,
        uint32_t               n_comp,
        uint32_t               window,
        uint32_t               ratio,
        uint32_t               n_head,
        uint32_t               head_dim) {
    /* The vector sibling below only handles the small square case; every
     * rectangular (TP row-split) problem goes through the long path. */
    if (n_tokens >= 20 || q_row0 != 0 || n_q != n_tokens) {
        return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(cbp,
                                                                                       heads,
                                                                                       sinks_buf,
                                                                                       sinks_offset,
                                                                                       q,
                                                                                       raw_kv,
                                                                                       comp_kv,
                                                                                       comp_kv_f16,
                                                                                       comp_mask,
                                                                                       use_comp_mask,
                                                                                       q_row0,
                                                                                       n_q,
                                                                                       n_tokens,
                                                                                       n_comp,
                                                                                       window,
                                                                                       ratio,
                                                                                       n_head,
                                                                                       head_dim);
    }
    return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(cbp,
                                                                           heads,
                                                                           sinks_buf,
                                                                           sinks_offset,
                                                                           q,
                                                                           raw_kv,
                                                                           comp_kv,
                                                                           comp_kv_f16,
                                                                           comp_mask,
                                                                           use_comp_mask,
                                                                           n_tokens,
                                                                           n_comp,
                                                                           window,
                                                                           ratio,
                                                                           n_head,
                                                                           head_dim);
}

/* Raw prefill FlashAttention over a rectangular problem: q holds n_q query
 * rows that correspond to token positions [q_row0, q_row0 + n_q) of the
 * chunk, raw_kv holds all n_kv key rows, and heads receives one output row
 * per query row.  The classic square prefill is q_row0 == 0, n_q == n_kv. */
static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(
        id<MTLCommandBuffer> __strong *cbp,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t               q_row0,
        uint32_t               n_q,
        uint32_t               n_kv,
        uint32_t               window,
        uint32_t               n_head,
        uint32_t               head_dim) {
    if (!cbp || !*cbp) return 0;
    id<MTLCommandBuffer> cb = *cbp;
    if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) {
        return 0;
    }

    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)n_kv * head_dim * sizeof(float);
    if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        ds4_gpu_tensor_bytes(heads) < q_bytes) {
        fprintf(stderr, "ds4: Metal prefill raw DS4 non-vector FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t nqptg = 8;
    const uint32_t ncpsg = 64;
    const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
    const bool has_kvpad = (n_kv % ncpsg) != 0;
    const bool bc_mask = (n_q % nqptg) != 0;
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_q * (NSUInteger)n_kv * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_kv * row_bytes_f16;
    const NSUInteger pad_bytes = has_kvpad
        ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t))
        : 1u;
    const NSUInteger nblk0 = ((NSUInteger)n_kv + ncpsg - 1u) / ncpsg;
    const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg;
    const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);

    bool mask_cache_created = false;
    ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache =
        q_row0 == 0u && n_q == n_kv
            ? ds4_gpu_get_zero_prefix_prefill_mask_cache(
                  DS4_GPU_PREFILL_MASK_CACHE_RAW,
                  n_kv,
                  0u,
                  n_kv,
                  window,
                  0u,
                  nqptg,
                  ncpsg,
                  has_kvpad,
                  bc_mask,
                  mask_bytes,
                  blk_bytes,
                  &mask_cache_created)
            : NULL;
    id<MTLBuffer> mask_buffer = mask_cache
        ? mask_cache->mask
        : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (mask_cache && mask_cache_created) {
        ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents],
                                      0u, n_kv, n_kv, window);
        mask_cache->valid = true;
    }
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv_f16") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        (!mask_cache &&
         !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
                                          &g_flash_attn_blk_bytes,
                                          blk_bytes,
                                          "ds4_flash_attn_blk"))) {
        return 0;
    }
    id<MTLBuffer> blk_buffer = mask_cache
        ? mask_cache->blk
        : g_flash_attn_blk_buffer;

    const bool flash_stage_profile =
        getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
    double flash_stage_t0 = 0.0;
    if (flash_stage_profile) {
        if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
            return 0;
        }
        int profile_owned = 0;
        cb = ds4_gpu_command_buffer(&profile_owned);
        if (!cb || profile_owned) return 0;
        *cbp = cb;
        flash_stage_t0 = ds4_gpu_now_ms();
    }
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
        if (flash_stage_profile) { \
            if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
                    "raw_nonvec", (name), n_q, 0, n_kv, \
                    n_head, head_dim, window, 0, &flash_stage_t0)) { \
                return 0; \
            } \
            cb = *cbp; \
        } \
    } while (0)

    if (!mask_cache) {
        ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents],
                                      q_row0, n_q, n_kv, window);
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");

    id<MTLComputePipelineState> pad_pipeline = nil;
    if (has_kvpad) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> blk_pipeline =
        ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
    id<MTLComputePipelineState> attn_pipeline =
        ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
                                          true, true, false, false, has_kvpad, bc_mask,
                                          (int32_t)head_dim,
                                          (int32_t)head_dim,
                                          (int32_t)nsg);
    if (!blk_pipeline || !attn_pipeline) return 0;

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         rawbuf,
                                         ds4_gpu_tensor_offset(raw_kv),
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_kv * head_dim)) {
        return 0;
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");

    if (has_kvpad) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_kv,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_kv * row_bytes_f16,
            .nb13 = (uint64_t)n_kv * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_kv * row_bytes_f16,
            .nb23 = (uint64_t)n_kv * row_bytes_f16,
            .ne31 = (int32_t)n_q,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)n_kv * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
    }

    ds4_gpu_flash_attn_blk_args blk_args = {
        .ne01 = (int32_t)n_q,
        .ne30 = (int32_t)n_kv,
        .ne31 = (int32_t)n_q,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_kv * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
    };

    id<MTLComputeCommandEncoder> enc = nil;
    if (!mask_cache || !mask_cache->blk_ready) {
        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:blk_pipeline];
        [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
        [enc setBuffer:mask_buffer offset:0 atIndex:1];
        [enc setBuffer:blk_buffer offset:0 atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        if (mask_cache) mask_cache->blk_ready = true;
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map");

    ds4_gpu_flash_attn_vec_args args = {
        .ne01 = (int32_t)n_q,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_q * n_head * row_bytes,
        .ne11 = (int32_t)n_kv,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_kv * row_bytes_f16,
        .nb13 = (uint64_t)n_kv * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_kv * row_bytes_f16,
        .nb23 = (uint64_t)n_kv * row_bytes_f16,
        .ne31 = (int32_t)n_q,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_kv * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = (int32_t)n_q,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
    const NSUInteger shared_elems = (NSUInteger)nqptg *
        ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:attn_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:blk_buffer offset:0 atIndex:7];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention");

#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
    return 1;
}

static int ds4_gpu_encode_flash_attention_prefill_raw_heads(
        id<MTLCommandBuffer> __strong *cbp,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t               q_row0,
        uint32_t               n_q,
        uint32_t               n_kv,
        uint32_t               window,
        uint32_t               n_head,
        uint32_t               head_dim) {
    if (!cbp || !*cbp) return 0;
    id<MTLCommandBuffer> cb = *cbp;
    if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) {
        return 0;
    }
    /* The vector sibling below only handles the small square case; every
     * rectangular (TP row-split) problem goes through the non-vector path. */
    if (n_kv >= 20 || q_row0 != 0 || n_q != n_kv) {
        return ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(cbp,
                                                                         heads,
                                                                         sinks_buf,
                                                                         sinks_offset,
                                                                         q,
                                                                         raw_kv,
                                                                         q_row0,
                                                                         n_q,
                                                                         n_kv,
                                                                         window,
                                                                         n_head,
                                                                         head_dim);
    }
    const uint32_t n_tokens = n_q;

    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
    if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        ds4_gpu_tensor_bytes(heads) < q_bytes) {
        fprintf(stderr, "ds4: Metal prefill raw DS4 FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t ncpsg = 32;
    const uint32_t nwg = 32;
    const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_tokens, nwg, ncpsg);
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)n_tokens * sizeof(uint16_t);
    const NSUInteger kv_f16_offset = 0;
    const NSUInteger kv_f16_bytes = (NSUInteger)n_tokens * row_bytes_f16;
    const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 +
                                 (NSUInteger)ncpsg * (NSUInteger)n_tokens * sizeof(uint16_t);
    const NSUInteger nrows = (NSUInteger)n_tokens * n_head;
    const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
                                 nrows * (2u * (NSUInteger)nwg) * sizeof(float);

    id<MTLBuffer> mask_buffer =
        ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_f16_bytes,
                                         "ds4_flash_attn_kv_f16") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
                                         &g_flash_attn_tmp_bytes,
                                         tmp_bytes,
                                         "ds4_flash_attn_tmp")) {
        return 0;
    }

    const bool flash_stage_profile =
        getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
    double flash_stage_t0 = 0.0;
    if (flash_stage_profile) {
        if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
            return 0;
        }
        int profile_owned = 0;
        cb = ds4_gpu_command_buffer(&profile_owned);
        if (!cb || profile_owned) return 0;
        *cbp = cb;
        flash_stage_t0 = ds4_gpu_now_ms();
    }
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
        if (flash_stage_profile) { \
            if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
                    "raw_vec", (name), n_tokens, 0, n_tokens, \
                    n_head, head_dim, window, 0, &flash_stage_t0)) { \
                return 0; \
            } \
            cb = *cbp; \
        } \
    } while (0)

    ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], 0, n_tokens, n_tokens, window);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");

    id<MTLComputePipelineState> pad_pipeline = nil;
    if ((n_tokens % ncpsg) != 0) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> vec_pipeline =
        ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
                                              true, true, false, false, true,
                                              false,
                                              (int32_t)head_dim,
                                              (int32_t)head_dim,
                                              (int32_t)nsg,
                                              (int32_t)nwg);
    id<MTLComputePipelineState> reduce_pipeline =
        ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
    if (!vec_pipeline || !reduce_pipeline) return 0;

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         rawbuf,
                                         ds4_gpu_tensor_offset(raw_kv),
                                         g_flash_attn_kv_buffer,
                                         kv_f16_offset,
                                         n_tokens * head_dim)) {
        return 0;
    }
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");

    if ((n_tokens % ncpsg) != 0) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_tokens,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_tokens * row_bytes_f16,
            .nb13 = (uint64_t)n_tokens * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_tokens * row_bytes_f16,
            .nb23 = (uint64_t)n_tokens * row_bytes_f16,
            .ne31 = (int32_t)n_tokens,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)n_tokens * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
    }

    ds4_gpu_flash_attn_vec_args vec_args = {
        .ne01 = (int32_t)n_tokens,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_tokens * n_head * row_bytes,
        .ne11 = (int32_t)n_tokens,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_tokens * row_bytes_f16,
        .nb13 = (uint64_t)n_tokens * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_tokens * row_bytes_f16,
        .nb23 = (uint64_t)n_tokens * row_bytes_f16,
        .ne31 = (int32_t)n_tokens,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_tokens * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = (int32_t)n_tokens,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
                                     4u * ncpsg +
                                     2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:vec_pipeline];
    [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec");

    ds4_gpu_flash_attn_reduce_args reduce_args = {
        .nrows = (int32_t)nrows,
    };
    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:reduce_pipeline];
    [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce");

#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
    return 1;
}

static int ds4_gpu_encode_flash_attention_gathered_heads(
        id<MTLCommandBuffer>   cb,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t               n_raw,
        uint32_t               raw_cap,
        uint32_t               raw_start,
        const ds4_gpu_tensor *comp_kv,
        uint32_t               comp_kv_f16,
        uint32_t               n_comp,
        const ds4_gpu_tensor *comp_mask,
        uint32_t               use_mask,
        uint32_t               n_head,
        uint32_t               head_dim) {
    const uint32_t n_keys = n_raw + n_comp;
    if (head_dim != 512 || n_head == 0 || n_raw == 0 || n_keys == 0 ||
        raw_cap < n_raw || n_keys < n_raw) {
        return 0;
    }

    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : nil;
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    id<MTLBuffer> maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : nil;
    const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
    const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
                                (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
    const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u;
    if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
        (n_comp && !compbuf) ||
        (use_mask && !maskbuf) ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
        ds4_gpu_tensor_bytes(heads) < q_bytes ||
        (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes)) {
        fprintf(stderr, "ds4: Metal gathered DS4 FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t ncpsg = 32;
    const uint32_t nwg = 32;
    const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg);
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_keys * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
    const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 +
                                 (NSUInteger)ncpsg * sizeof(uint16_t);
    const NSUInteger nrows = (NSUInteger)n_head;
    const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
                                 nrows * (2u * (NSUInteger)nwg) * sizeof(float);

    const bool use_persistent_zero_mask =
        use_mask == 0u &&
        (ds4_gpu_device_name_contains("M3") ||
         getenv("DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK") != NULL) &&
        getenv("DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK") == NULL;

    if (!(use_persistent_zero_mask
              ? ds4_gpu_ensure_zero_attention_mask(mask_bytes)
              : ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer,
                                                &g_flash_attn_mask_bytes,
                                                mask_bytes,
                                                "ds4_flash_attn_mask")) ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
                                         &g_flash_attn_tmp_bytes,
                                         tmp_bytes,
                                         "ds4_flash_attn_tmp")) {
        return 0;
    }
    id<MTLBuffer> flash_mask_buffer = use_persistent_zero_mask
        ? g_flash_attn_zero_mask_buffer
        : g_flash_attn_mask_buffer;

    const bool has_kvpad = (n_keys % ncpsg) != 0;
    const bool use_shared_kvpad =
        has_kvpad &&
        (ds4_gpu_device_name_contains("M3") ||
         getenv("DS4_METAL_ENABLE_SHARED_KV_PAD") != NULL) &&
        getenv("DS4_METAL_DISABLE_M3_SHARED_KV_PAD") == NULL;
    id<MTLComputePipelineState> vec_pipeline =
        ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
                                              true, true, false, false, has_kvpad,
                                              use_shared_kvpad,
                                              (int32_t)head_dim,
                                              (int32_t)head_dim,
                                              (int32_t)nsg,
                                              (int32_t)nwg);
    id<MTLComputePipelineState> reduce_pipeline =
        ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
    if (!vec_pipeline || !reduce_pipeline) return 0;

    if (!use_persistent_zero_mask &&
        !ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) {
        return 0;
    }
    if (use_mask && n_comp &&
        !ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         maskbuf,
                                         ds4_gpu_tensor_offset(comp_mask),
                                         flash_mask_buffer,
                                         (NSUInteger)n_raw * sizeof(uint16_t),
                                         n_comp)) {
        return 0;
    }

    bool pad_fused = false;
    if (!ds4_gpu_encode_flash_kv_stage_f16(
            cb,
            rawbuf,
            ds4_gpu_tensor_offset(raw_kv),
            raw_cap,
            raw_start,
            n_raw,
            compbuf,
            ds4_gpu_tensor_offset(comp_kv),
            comp_kv_f16 != 0,
            n_comp,
            head_dim,
            g_flash_attn_kv_buffer,
            0,
            flash_mask_buffer,
            0,
            g_flash_attn_pad_buffer,
            0,
            has_kvpad,
            use_shared_kvpad,
            &pad_fused)) {
        return 0;
    }

    id<MTLComputePipelineState> pad_pipeline = nil;
    if (has_kvpad && !pad_fused) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }

    if (has_kvpad && !pad_fused) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_keys,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_keys * row_bytes_f16,
            .nb13 = (uint64_t)n_keys * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_keys * row_bytes_f16,
            .nb23 = (uint64_t)n_keys * row_bytes_f16,
            .ne31 = 1,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = mask_bytes,
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:flash_mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
    }

    ds4_gpu_flash_attn_vec_args vec_args = {
        .ne01 = 1,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_head * row_bytes,
        .ne11 = (int32_t)n_keys,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_keys * row_bytes_f16,
        .nb13 = (uint64_t)n_keys * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_keys * row_bytes_f16,
        .nb23 = (uint64_t)n_keys * row_bytes_f16,
        .ne31 = 1,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = mask_bytes,
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = 1,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
                                     4u * ncpsg +
                                     2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:vec_pipeline];
    [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:flash_mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    ds4_gpu_flash_attn_reduce_args reduce_args = {
        .nrows = (int32_t)nrows,
    };
    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:reduce_pipeline];
    [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
    [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads(
        id<MTLCommandBuffer>   cb,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t               n_tokens,
        uint32_t               pos0,
        uint32_t               n_raw,
        uint32_t               raw_cap,
        uint32_t               raw_start,
        uint32_t               window,
        uint32_t               n_head,
        uint32_t               head_dim,
        bool                   noncausal) {
    if (head_dim != 512 || n_head == 0 || n_tokens == 0 ||
        n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) {
        return 0;
    }

    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
    if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        ds4_gpu_tensor_bytes(heads) < q_bytes) {
        fprintf(stderr, "ds4: Metal decode raw batch FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t nqptg = 8;
    const uint32_t ncpsg = 64;
    const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
    const bool has_kvpad = (n_raw % ncpsg) != 0;
    const bool bc_mask = (n_tokens % nqptg) != 0;
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_raw * (NSUInteger)n_tokens * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16;
    const NSUInteger pad_bytes = has_kvpad
        ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t))
        : 1u;
    const NSUInteger nblk0 = ((NSUInteger)n_raw + ncpsg - 1u) / ncpsg;
    const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg;
    const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);

    id<MTLBuffer> mask_buffer =
        ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv_f16") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
                                         &g_flash_attn_blk_bytes,
                                         blk_bytes,
                                         "ds4_flash_attn_blk")) {
        return 0;
    }

    if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              raw_cap,
                                              raw_start,
                                              n_raw,
                                              head_dim,
                                              g_flash_attn_kv_buffer,
                                              0)) {
        return 0;
    }

    if (noncausal) {
        memset([mask_buffer contents], 0, mask_bytes);
    } else {
        ds4_gpu_fill_raw_decode_batch_mask((uint16_t *)[mask_buffer contents],
                                             n_tokens,
                                             n_raw,
                                             pos0,
                                             window);
    }

    id<MTLComputePipelineState> pad_pipeline = nil;
    if (has_kvpad) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> blk_pipeline =
        ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
    id<MTLComputePipelineState> attn_pipeline =
        ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
                                          true, true, false, false, has_kvpad, bc_mask,
                                          (int32_t)head_dim,
                                          (int32_t)head_dim,
                                          (int32_t)nsg);
    if (!blk_pipeline || !attn_pipeline) return 0;

    if (has_kvpad) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_raw,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_raw * row_bytes_f16,
            .nb13 = (uint64_t)n_raw * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_raw * row_bytes_f16,
            .nb23 = (uint64_t)n_raw * row_bytes_f16,
            .ne31 = (int32_t)n_tokens,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)n_raw * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
    }

    ds4_gpu_flash_attn_blk_args blk_args = {
        .ne01 = (int32_t)n_tokens,
        .ne30 = (int32_t)n_raw,
        .ne31 = (int32_t)n_tokens,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_raw * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
    };

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:blk_pipeline];
    [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
    [enc setBuffer:mask_buffer offset:0 atIndex:1];
    [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    ds4_gpu_flash_attn_vec_args args = {
        .ne01 = (int32_t)n_tokens,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_tokens * n_head * row_bytes,
        .ne11 = (int32_t)n_raw,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_raw * row_bytes_f16,
        .nb13 = (uint64_t)n_raw * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_raw * row_bytes_f16,
        .nb23 = (uint64_t)n_raw * row_bytes_f16,
        .ne31 = (int32_t)n_tokens,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_raw * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = (int32_t)n_tokens,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
    const NSUInteger shared_elems = (NSUInteger)nqptg *
        ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:attn_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
        id<MTLCommandBuffer>   cb,
        ds4_gpu_tensor      *heads,
        id<MTLBuffer>          sinks_buf,
        NSUInteger             sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t               comp_kv_f16,
        const ds4_gpu_tensor *comp_mask,
        uint32_t               use_comp_mask,
        uint32_t               n_tokens,
        uint32_t               pos0,
        uint32_t               n_raw,
        uint32_t               raw_cap,
        uint32_t               raw_start,
        uint32_t               n_comp,
        uint32_t               window,
        uint32_t               ratio,
        uint32_t               n_head,
        uint32_t               head_dim) {
    if (n_comp == 0) {
        return ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
                                                                       heads,
                                                                       sinks_buf,
                                                                       sinks_offset,
                                                                       q,
                                                                       raw_kv,
                                                                       n_tokens,
                                                                       pos0,
                                                                       n_raw,
                                                                       raw_cap,
                                                                       raw_start,
                                                                       window,
                                                                       n_head,
                                                                       head_dim,
                                                                       false);
    }
    if (head_dim != 512 || n_head == 0 || n_tokens == 0 ||
        n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
        ratio == 0 || !comp_kv || (use_comp_mask && !comp_mask)) {
        return 0;
    }

    const uint32_t n_keys = n_raw + n_comp;
    id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
    id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
    id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_kv);
    id<MTLBuffer> maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
    id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
    const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
    const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
    const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
                                (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
    const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u;
    if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf ||
        ds4_gpu_tensor_bytes(q) < q_bytes ||
        ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
        ds4_gpu_tensor_bytes(comp_kv) < comp_bytes ||
        (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
        ds4_gpu_tensor_bytes(heads) < q_bytes) {
        fprintf(stderr, "ds4: Metal decode mixed batch FlashAttention received undersized buffers\n");
        return 0;
    }

    const uint32_t nqptg = 8;
    const uint32_t ncpsg = 64;
    const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
    const bool has_kvpad = (n_keys % ncpsg) != 0;
    const bool bc_mask = (n_tokens % nqptg) != 0;
    const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
    const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
    const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
    const NSUInteger pad_bytes = has_kvpad
        ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t))
        : 1u;
    const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg;
    const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg;
    const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);

    id<MTLBuffer> mask_buffer =
        ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
    if (!mask_buffer ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                         &g_flash_attn_kv_bytes,
                                         kv_bytes,
                                         "ds4_flash_attn_kv_f16") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                         &g_flash_attn_pad_bytes,
                                         pad_bytes,
                                         "ds4_flash_attn_pad") ||
        !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
                                         &g_flash_attn_blk_bytes,
                                         blk_bytes,
                                         "ds4_flash_attn_blk")) {
        return 0;
    }

    if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              raw_cap,
                                              raw_start,
                                              n_raw,
                                              head_dim,
                                              g_flash_attn_kv_buffer,
                                              0) ||
        !ds4_gpu_encode_copy_to_f16_1d(cb,
                                       compbuf,
                                       ds4_gpu_tensor_offset(comp_kv),
                                       comp_kv_f16 != 0,
                                       g_flash_attn_kv_buffer,
                                       (NSUInteger)n_raw * row_bytes_f16,
                                       n_comp * head_dim)) {
        return 0;
    }

    ds4_gpu_fill_mixed_decode_batch_mask((uint16_t *)[mask_buffer contents],
                                           n_tokens,
                                           n_raw,
                                           n_comp,
                                           pos0,
                                           window,
                                           ratio);
    if (use_comp_mask) {
        if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
                                             maskbuf,
                                             ds4_gpu_tensor_offset(comp_mask),
                                             mask_buffer,
                                             (NSUInteger)n_raw * sizeof(uint16_t),
                                             n_comp,
                                             n_tokens,
                                             (uint64_t)n_comp * sizeof(float),
                                             (uint64_t)n_keys * sizeof(uint16_t))) {
            return 0;
        }
    }

    id<MTLComputePipelineState> pad_pipeline = nil;
    if (has_kvpad) {
        pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
        if (!pad_pipeline) return 0;
    }
    id<MTLComputePipelineState> blk_pipeline =
        ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
    id<MTLComputePipelineState> attn_pipeline =
        ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
                                          true, true, false, false, has_kvpad, bc_mask,
                                          (int32_t)head_dim,
                                          (int32_t)head_dim,
                                          (int32_t)nsg);
    if (!blk_pipeline || !attn_pipeline) return 0;

    if (has_kvpad) {
        ds4_gpu_flash_attn_pad_args pad_args = {
            .ne11 = (int32_t)n_keys,
            .ne_12_2 = 1,
            .ne_12_3 = 1,
            .nb11 = row_bytes_f16,
            .nb12 = (uint64_t)n_keys * row_bytes_f16,
            .nb13 = (uint64_t)n_keys * row_bytes_f16,
            .nb21 = row_bytes_f16,
            .nb22 = (uint64_t)n_keys * row_bytes_f16,
            .nb23 = (uint64_t)n_keys * row_bytes_f16,
            .ne31 = (int32_t)n_tokens,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pad_pipeline];
        [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
        [enc setBuffer:mask_buffer offset:0 atIndex:3];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
    }

    ds4_gpu_flash_attn_blk_args blk_args = {
        .ne01 = (int32_t)n_tokens,
        .ne30 = (int32_t)n_keys,
        .ne31 = (int32_t)n_tokens,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
    };

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:blk_pipeline];
    [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
    [enc setBuffer:mask_buffer offset:0 atIndex:1];
    [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    ds4_gpu_flash_attn_vec_args args = {
        .ne01 = (int32_t)n_tokens,
        .ne02 = (int32_t)n_head,
        .ne03 = 1,
        .nb01 = (uint64_t)n_head * row_bytes,
        .nb02 = row_bytes,
        .nb03 = (uint64_t)n_tokens * n_head * row_bytes,
        .ne11 = (int32_t)n_keys,
        .ne_12_2 = 1,
        .ne_12_3 = 1,
        .ns10 = (int32_t)head_dim,
        .nb11 = row_bytes_f16,
        .nb12 = (uint64_t)n_keys * row_bytes_f16,
        .nb13 = (uint64_t)n_keys * row_bytes_f16,
        .ns20 = (int32_t)head_dim,
        .nb21 = row_bytes_f16,
        .nb22 = (uint64_t)n_keys * row_bytes_f16,
        .nb23 = (uint64_t)n_keys * row_bytes_f16,
        .ne31 = (int32_t)n_tokens,
        .ne32 = 1,
        .ne33 = 1,
        .nb31 = (uint64_t)n_keys * sizeof(uint16_t),
        .nb32 = mask_bytes,
        .nb33 = mask_bytes,
        .ne1 = (int32_t)n_head,
        .ne2 = (int32_t)n_tokens,
        .ne3 = 1,
        .scale = 1.0f / sqrtf((float)head_dim),
        .max_bias = 0.0f,
        .m0 = 0.0f,
        .m1 = 0.0f,
        .n_head_log2 = 0,
        .logit_softcap = 0.0f,
    };

    const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
    const NSUInteger shared_elems = (NSUInteger)nqptg *
        ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
    const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:attn_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
    [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
    [enc setBuffer:mask_buffer offset:0 atIndex:4];
    [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
    [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
    [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7];
    [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
    [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

int ds4_gpu_attention_prefill_raw_heads_range_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t                q_row0,
        uint32_t                n_q,
        uint32_t                n_kv,
        uint32_t                window,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_kv == 0) return 0;

    @autoreleasepool {
        if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             (uint64_t)n_head * sizeof(float),
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_prefill_raw_heads(&cb,
                                                                heads,
                                                                sinks_buf,
                                                                (NSUInteger)sinks_inner,
                                                                q,
                                                                raw_kv,
                                                                q_row0,
                                                                n_q,
                                                                n_kv,
                                                                window,
                                                                n_head,
                                                                head_dim)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill raw attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_prefill_raw_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t                n_tokens,
        uint32_t                window,
        uint32_t                n_head,
        uint32_t                head_dim) {
    return ds4_gpu_attention_prefill_raw_heads_range_tensor(heads,
                                                              model_map,
                                                              model_size,
                                                              sinks_offset,
                                                              q,
                                                              raw_kv,
                                                              0,
                                                              n_tokens,
                                                              n_tokens,
                                                              window,
                                                              n_head,
                                                              head_dim);
}

int ds4_gpu_attention_decode_raw_batch_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t                n_tokens,
        uint32_t                pos0,
        uint32_t                n_raw,
        uint32_t                raw_cap,
        uint32_t                raw_start,
        uint32_t                window,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 ||
        n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) {
        return 0;
    }

    @autoreleasepool {
        if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             (uint64_t)n_head * sizeof(float),
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
                                                                     heads,
                                                                     sinks_buf,
                                                                     (NSUInteger)sinks_inner,
                                                                     q,
                                                                     raw_kv,
                                                                     n_tokens,
                                                                     pos0,
                                                                     n_raw,
                                                                     raw_cap,
                                                                     raw_start,
                                                                     window,
                                                                     n_head,
                                                                     head_dim,
                                                                     false)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode raw batch attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t                n_tokens,
        uint32_t                n_raw,
        uint32_t                raw_cap,
        uint32_t                raw_start,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map ||
        n_tokens == 0 || n_raw == 0 || raw_cap < n_raw ||
        raw_start >= raw_cap || n_head == 0 || head_dim == 0) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float);
        if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal noncausal attention sinks range is outside the mapped model\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             sink_bytes,
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
                                                                     heads,
                                                                     sinks_buf,
                                                                     (NSUInteger)sinks_inner,
                                                                     q,
                                                                     raw_kv,
                                                                     n_tokens,
                                                                     0,
                                                                     n_raw,
                                                                     raw_cap,
                                                                     raw_start,
                                                                     0,
                                                                     n_head,
                                                                     head_dim,
                                                                     true)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph noncausal raw batch attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_decode_mixed_batch_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        const ds4_gpu_tensor *comp_mask,
        uint32_t                use_comp_mask,
        uint32_t                n_tokens,
        uint32_t                pos0,
        uint32_t                n_raw,
        uint32_t                raw_cap,
        uint32_t                raw_start,
        uint32_t                n_comp,
        uint32_t                window,
        uint32_t                ratio,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 ||
        n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
        ratio == 0 || (n_comp != 0 && !comp_kv) ||
        (use_comp_mask != 0 && !comp_mask)) {
        return 0;
    }

    @autoreleasepool {
        if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             (uint64_t)n_head * sizeof(float),
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(cb,
                                                                       heads,
                                                                       sinks_buf,
                                                                       (NSUInteger)sinks_inner,
                                                                       q,
                                                                       raw_kv,
                                                                       comp_kv,
                                                                       comp_kv_f16,
                                                                       comp_mask,
                                                                       use_comp_mask,
                                                                       n_tokens,
                                                                       pos0,
                                                                       n_raw,
                                                                       raw_cap,
                                                                       raw_start,
                                                                       n_comp,
                                                                       window,
                                                                       ratio,
                                                                       n_head,
                                                                       head_dim)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode mixed batch attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        const ds4_gpu_tensor *topk,
        uint32_t                n_tokens,
        uint32_t                pos0,
        uint32_t                n_raw,
        uint32_t                raw_cap,
        uint32_t                raw_start,
        uint32_t                n_comp,
        uint32_t                top_k,
        uint32_t                window,
        uint32_t                ratio,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !model_map || !q || !raw_kv || !comp_kv || !topk ||
        n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
        n_comp == 0 || top_k == 0 || top_k > n_comp || (top_k & (top_k - 1u)) != 0 ||
        ratio == 0 || n_head == 0 || head_dim != 512) {
        return 0;
    }

    @autoreleasepool {
        if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal indexed attention sinks range is outside the mapped model\n");
            return 0;
        }

        const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
        const uint64_t row_bytes_f16 = (uint64_t)head_dim * sizeof(uint16_t);
        const uint64_t q_bytes = (uint64_t)n_tokens * n_head * row_bytes;
        const uint64_t raw_bytes = (uint64_t)raw_cap * row_bytes;
        const uint64_t comp_bytes = (uint64_t)n_comp * (comp_kv_f16 ? row_bytes_f16 : row_bytes);
        const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
        id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_kv);
        id<MTLBuffer> topkbuf = ds4_gpu_tensor_buffer(topk);
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        if (!qbuf || !rawbuf || !compbuf || !topkbuf || !headsbuf ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
            ds4_gpu_tensor_bytes(comp_kv) < comp_bytes ||
            ds4_gpu_tensor_bytes(topk) < topk_bytes ||
            ds4_gpu_tensor_bytes(heads) < q_bytes) {
            fprintf(stderr, "ds4: Metal indexed mixed attention received undersized buffers\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             (uint64_t)n_head * sizeof(float),
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        id<MTLComputePipelineState> sort_pipeline =
            ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline,
                                    "kernel_dsv4_sort_i32_rows_asc");
        const bool decode_one_token = n_tokens == 1u;
        id<MTLComputePipelineState> attn_pipeline =
            decode_one_token ?
            ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline,
                                   "kernel_dsv4_indexed_mixed_attention_heads8_rb16") :
            ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline,
                                   "kernel_dsv4_indexed_mixed_attention_heads8");
        if (!sort_pipeline || !attn_pipeline) return 0;
        if ((NSUInteger)top_k > sort_pipeline.maxTotalThreadsPerThreadgroup) {
            fprintf(stderr, "ds4: Metal indexed attention top-k exceeds sort threadgroup limit\n");
            return 0;
        }
        /*
         * Fast decode attends to the same full top-k compressed rows but keeps
         * them in score order, avoiding a chronological sort dispatch.
         * --quality restores the sorted order for stricter reproducibility.
         */
        const bool skip_decode_sort = !g_quality_mode && decode_one_token;
        if (!skip_decode_sort &&
            !ds4_gpu_ensure_scratch_buffer(&g_indexed_topk_buffer,
                                             &g_indexed_topk_bytes,
                                             (NSUInteger)topk_bytes,
                                             "ds4_indexed_topk_sorted")) {
            return 0;
        }

        ds4_gpu_dsv4_topk_mask_args sort_args = {
            .ne00 = (int64_t)top_k,
            .ne01 = (int64_t)n_tokens,
            .nb00 = sizeof(int32_t),
            .nb01 = (uint64_t)top_k * sizeof(int32_t),
            .ne0 = (int64_t)top_k,
            .ne1 = (int64_t)n_tokens,
            .nb0 = sizeof(int32_t),
            .nb1 = (uint64_t)top_k * sizeof(int32_t),
        };
        ds4_gpu_dsv4_indexed_attention_args attn_args = {
            .n_tokens = n_tokens,
            .n_head = n_head,
            .n_raw = n_raw,
            .raw_cap = raw_cap,
            .raw_start = raw_start,
            .n_comp = n_comp,
            .top_k = top_k,
            .pos0 = pos0,
            .window = window,
            .ratio = ratio,
            .comp_kv_f16 = comp_kv_f16 ? 1u : 0u,
            .pad0 = 0,
            .q_token_stride = (uint64_t)n_head * row_bytes,
            .q_head_stride = row_bytes,
            .raw_row_stride = row_bytes,
            .comp_row_stride = comp_kv_f16 ? row_bytes_f16 : row_bytes,
            .topk_token_stride = (uint64_t)top_k * sizeof(int32_t),
            .dst_token_stride = (uint64_t)n_head * row_bytes,
            .dst_head_stride = row_bytes,
            .scale = 1.0f / sqrtf((float)head_dim),
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = nil;
        if (!skip_decode_sort) {
            enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:sort_pipeline];
            [enc setBytes:&sort_args length:sizeof(sort_args) atIndex:0];
            [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1];
            [enc setBuffer:g_indexed_topk_buffer offset:0 atIndex:2];
            [enc setThreadgroupMemoryLength:(NSUInteger)top_k * sizeof(int32_t) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(top_k, 1, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);
        }

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:attn_pipeline];
        [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2];
        [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3];
        [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer
              offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0
             atIndex:4];
        [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6];
        [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) *
                                        128u * 4u * sizeof(uint16_t)
                                atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, ((NSUInteger)n_head + 7u) / 8u, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph indexed mixed attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        uint32_t                q_row0,
        uint32_t                n_q,
        uint32_t                n_tokens,
        uint32_t                n_comp,
        uint32_t                window,
        uint32_t                ratio,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_tokens == 0 ||
        ratio == 0 || (n_comp != 0 && !comp_kv)) {
        return 0;
    }

    @autoreleasepool {
        if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             (uint64_t)n_head * sizeof(float),
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb,
                                                                                heads,
                                                                                sinks_buf,
                                                                                (NSUInteger)sinks_inner,
                                                                                q,
                                                                                raw_kv,
                                                                                comp_kv,
                                                                                comp_kv_f16,
                                                                                NULL,
                                                                                0,
                                                                                q_row0,
                                                                                n_q,
                                                                                n_tokens,
                                                                                n_comp,
                                                                                window,
                                                                                ratio,
                                                                                n_head,
                                                                                head_dim)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill static mixed attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_prefill_static_mixed_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        uint32_t                n_tokens,
        uint32_t                n_comp,
        uint32_t                window,
        uint32_t                ratio,
        uint32_t                n_head,
        uint32_t                head_dim) {
    return ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(heads,
                                                                       model_map,
                                                                       model_size,
                                                                       sinks_offset,
                                                                       q,
                                                                       raw_kv,
                                                                       comp_kv,
                                                                       comp_kv_f16,
                                                                       0,
                                                                       n_tokens,
                                                                       n_tokens,
                                                                       n_comp,
                                                                       window,
                                                                       ratio,
                                                                       n_head,
                                                                       head_dim);
}

int ds4_gpu_attention_prefill_masked_mixed_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        const ds4_gpu_tensor *comp_mask,
        uint32_t                n_tokens,
        uint32_t                n_comp,
        uint32_t                window,
        uint32_t                ratio,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !comp_kv || !comp_mask || !model_map ||
        n_tokens == 0 || n_comp == 0 || ratio == 0) {
        return 0;
    }

    @autoreleasepool {
        if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                             sinks_offset,
                                                             (uint64_t)n_head * sizeof(float),
                                                             &sinks_inner);
        if (!sinks_buf) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb,
                                                                                heads,
                                                                                sinks_buf,
                                                                                (NSUInteger)sinks_inner,
                                                                                q,
                                                                                raw_kv,
                                                                                comp_kv,
                                                                                comp_kv_f16,
                                                                                comp_mask,
                                                                                1,
                                                                                0,
                                                                                n_tokens,
                                                                                n_tokens,
                                                                                n_comp,
                                                                                window,
                                                                                ratio,
                                                                                n_head,
                                                                                head_dim)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill masked mixed attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_attention_decode_heads_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t                n_raw,
        uint32_t                raw_cap,
        uint32_t                raw_start,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        uint32_t                n_comp,
        const ds4_gpu_tensor *comp_mask,
        uint32_t                use_mask,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !model_map || !q || !raw_kv ||
        n_raw == 0 || n_head == 0 || head_dim == 0 ||
        raw_cap < n_raw || raw_start >= raw_cap ||
        n_raw > UINT32_MAX - n_comp || n_raw + n_comp > 8192u ||
        (n_comp != 0 && !comp_kv) ||
        (use_mask != 0 && !comp_mask)) {
        return 0;
    }

    @autoreleasepool {
        const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
        const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
        const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
                                    (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
        const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float);
        if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) {
            fprintf(stderr, "ds4: Metal graph attention heads sink range is outside the mapped model\n");
            return 0;
        }

        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
        id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf;
        id<MTLBuffer> maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u;
        if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
            (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
            (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
            ds4_gpu_tensor_bytes(heads) < q_bytes) {
            fprintf(stderr, "ds4: Metal graph attention heads received undersized buffers\n");
            return 0;
        }

        uint64_t sinks_inner = 0;
        id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset, sink_bytes, &sinks_inner);
        if (!sinks_buf) return 0;

        if (n_comp == 0) {
            int owned = 0;
            id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;

            if (!ds4_gpu_encode_flash_attention_raw_heads(cb,
                                                            heads,
                                                            sinks_buf,
                                                            (NSUInteger)sinks_inner,
                                                            q,
                                                            raw_kv,
                                                            n_raw,
                                                            raw_cap,
                                                            raw_start,
                                                            n_head,
                                                            head_dim)) {
                return 0;
            }

            if (!ds4_gpu_finish_command_buffer(cb, owned, "graph raw attention heads")) return 0;
            return 1;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!ds4_gpu_encode_flash_attention_gathered_heads(cb,
                                                             heads,
                                                             sinks_buf,
                                                             (NSUInteger)sinks_inner,
                                                             q,
                                                             raw_kv,
                                                             n_raw,
                                                             raw_cap,
                                                             raw_start,
                                                             comp_kv,
                                                             comp_kv_f16,
                                                             n_comp,
                                                             comp_mask,
                                                             use_mask,
                                                             n_head,
                                                             head_dim)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph attention heads")) return 0;
    }

    return 1;
}

int ds4_gpu_swiglu_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *gate,
        const ds4_gpu_tensor *up,
        uint32_t                n,
        float                   clamp,
        float                   weight) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !gate || !up || n == 0) return 0;
    if (!isfinite(clamp) || clamp < 0.0f || !isfinite(weight)) return 0;

    @autoreleasepool {
        id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
        id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t bytes = (uint64_t)n * sizeof(float);
        if (!gatebuf || !upbuf || !outbuf ||
            ds4_gpu_tensor_bytes(gate) < bytes ||
            ds4_gpu_tensor_bytes(up) < bytes ||
            ds4_gpu_tensor_bytes(out) < bytes) {
            fprintf(stderr, "ds4: Metal SwiGLU received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glu_args args = {
            .ne00 = (int32_t)n,
            .nb01 = (uint64_t)n * sizeof(float),
            .ne10 = (int32_t)n,
            .nb11 = (uint64_t)n * sizeof(float),
            .ne0 = (int32_t)n,
            .nb1 = (uint64_t)n * sizeof(float),
            .i00 = 0,
            .i10 = 0,
            .alpha = weight,
            .limit = clamp,
        };
        NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > 256u) nth = 256u;
        if (nth > (NSUInteger)n) nth = (NSUInteger)n;
        if (nth == 0u) nth = 1u;
        const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_swiglu_flat_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:1];
        [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "SwiGLU")) return 0;
    }

    return 1;
}

int ds4_gpu_add_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *a,
        const ds4_gpu_tensor *b,
        uint32_t                n) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !a || !b || n == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
        id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t bytes = (uint64_t)n * sizeof(float);
        if (!abuf || !bbuf || !outbuf ||
            ds4_gpu_tensor_bytes(a) < bytes ||
            ds4_gpu_tensor_bytes(b) < bytes ||
            ds4_gpu_tensor_bytes(out) < bytes) {
            fprintf(stderr, "ds4: Metal tensor add received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_add_flat_args args = { .n = n };
        NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > 256u) nth = 256u;
        if (nth > (NSUInteger)n) nth = (NSUInteger)n;
        if (nth == 0u) nth = 1u;
        const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_add2_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
        [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add")) return 0;
    }

    return 1;
}

int ds4_gpu_add3_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *a,
        const ds4_gpu_tensor *b,
        const ds4_gpu_tensor *c,
        uint32_t                n) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !a || !b || !c || n == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
        id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
        id<MTLBuffer> cbuf = ds4_gpu_tensor_buffer(c);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t bytes = (uint64_t)n * sizeof(float);
        if (!abuf || !bbuf || !cbuf || !outbuf ||
            ds4_gpu_tensor_bytes(a) < bytes ||
            ds4_gpu_tensor_bytes(b) < bytes ||
            ds4_gpu_tensor_bytes(c) < bytes ||
            ds4_gpu_tensor_bytes(out) < bytes) {
            fprintf(stderr, "ds4: Metal tensor add3 received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_add_flat_args args = { .n = n };
        NSUInteger nth = g_add3_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > 256u) nth = 256u;
        if (nth > (NSUInteger)n) nth = (NSUInteger)n;
        if (nth == 0u) nth = 1u;
        const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_add3_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
        [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
        [enc setBuffer:cbuf offset:ds4_gpu_tensor_offset(c) atIndex:3];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add3")) return 0;
    }

    return 1;
}

typedef struct {
    uint32_t width;
    uint32_t rows;
    uint32_t layer;
    uint32_t n_threads;
    float    scale;
} ds4_gpu_directional_steering_project_args;

int ds4_gpu_directional_steering_project_tensor(
        ds4_gpu_tensor       *x,
        const ds4_gpu_tensor *directions,
        uint32_t                layer,
        uint32_t                width,
        uint32_t                rows,
        float                   scale) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0;

    @autoreleasepool {
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_directional_steering_project_f32");
        if (!pipeline) return 0;

        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> dbuf = ds4_gpu_tensor_buffer(directions);
        const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float);
        const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float);
        if (!xbuf || !dbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(directions) < dir_bytes) {
            fprintf(stderr, "ds4: Metal directional steering received undersized buffers\n");
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > 256u) nth = 256u;
        while (nth > width && nth > 1u) nth >>= 1;
        if (nth == 0) nth = 1;

        ds4_gpu_directional_steering_project_args args = {
            .width = width,
            .rows = rows,
            .layer = layer,
            .n_threads = (uint32_t)nth,
            .scale = scale,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc setBuffer:dbuf offset:ds4_gpu_tensor_offset(directions) atIndex:2];
        [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "directional steering")) return 0;
    }

    return 1;
}

static NSUInteger ds4_gpu_bin_threads(uint32_t width, id<MTLComputePipelineState> pipeline) {
    NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup;
    if (nth_max > 256u) nth_max = 256u;
    NSUInteger nth = 1u;
    while (2u * nth < (NSUInteger)width && nth < nth_max) nth *= 2u;
    return nth ? nth : 1u;
}

static int ds4_gpu_encode_unary_f32_rows(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        id<MTLBuffer>               src,
        NSUInteger                  src_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        uint32_t                    width,
        uint32_t                    rows,
        int                         c4,
        float                       min,
        float                       max) {
    if (!cb || !pipeline || !src || !dst || width == 0 || rows == 0) return 0;
    if (c4 && (width & 3u) != 0) return 0;

    ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, c4, 0.0f, 0.0f);
    args.min = min;
    args.max = max;

    NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup;
    if (nth_max > 256u) nth_max = 256u;
    NSUInteger nth = (NSUInteger)args.ne00;
    if (nth > nth_max) nth = nth_max;
    if (nth == 0) nth = 1u;
    const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01,
                                          (NSUInteger)args.ne02,
                                          (NSUInteger)args.ne03)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_bin_f32_rows(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_bin_args   *args,
        id<MTLBuffer>               a,
        NSUInteger                  a_off,
        id<MTLBuffer>               b,
        NSUInteger                  b_off,
        id<MTLBuffer>               out,
        NSUInteger                  out_off) {
    if (!cb || !pipeline || !args || !a || !b || !out || args->ne0 <= 0 || args->ne1 <= 0) {
        return 0;
    }

    const NSUInteger nth = ds4_gpu_bin_threads((uint32_t)args->ne0, pipeline);
    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:a offset:a_off atIndex:1];
    [enc setBuffer:b offset:b_off atIndex:2];
    [enc setBuffer:out offset:out_off atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)args->ne1,
                                          (NSUInteger)args->ne2,
                                          (NSUInteger)args->ne3)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static ds4_gpu_bin_args ds4_gpu_make_bin_rowwise_scalar_args(uint32_t width, uint32_t rows) {
    const uint64_t lhs_row_bytes = (uint64_t)width * sizeof(float);
    const uint64_t rhs_row_bytes = sizeof(float);
    return (ds4_gpu_bin_args) {
        .ne00 = (int32_t)width,
        .ne01 = (int32_t)rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = lhs_row_bytes,
        .nb02 = (uint64_t)rows * lhs_row_bytes,
        .nb03 = (uint64_t)rows * lhs_row_bytes,
        .ne10 = 1,
        .ne11 = (int32_t)rows,
        .ne12 = 1,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = rhs_row_bytes,
        .nb12 = (uint64_t)rows * rhs_row_bytes,
        .nb13 = (uint64_t)rows * rhs_row_bytes,
        .ne0 = (int32_t)width,
        .ne1 = (int32_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = lhs_row_bytes,
        .nb2 = (uint64_t)rows * lhs_row_bytes,
        .nb3 = (uint64_t)rows * lhs_row_bytes,
        .offs = 0,
        .o1 = { 0 },
    };
}

static ds4_gpu_mul_mv_id_args ds4_gpu_make_mul_mv_id_args(
        uint32_t src0_cols,
        uint32_t src0_rows,
        uint32_t src0_experts,
        uint64_t src0_row_bytes,
        uint64_t src0_expert_bytes,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens,
        uint32_t nr0) {
    const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float);
    const uint64_t src0_blocks = src0_cols / 256u;
    const uint64_t src0_block_bytes = src0_blocks ? src0_row_bytes / src0_blocks : 1u;
    return (ds4_gpu_mul_mv_id_args) {
        .nei0 = (int32_t)selected_experts,
        .nei1 = (int32_t)n_tokens,
        .nbi1 = (uint64_t)selected_experts * sizeof(int32_t),
        .ne00 = (int32_t)src0_cols,
        .ne01 = (int32_t)src0_rows,
        .ne02 = (int32_t)src0_experts,
        .nb00 = src0_block_bytes,
        .nb01 = src0_row_bytes,
        .nb02 = src0_expert_bytes,
        .ne10 = (int32_t)src0_cols,
        .ne11 = (int32_t)src1_expert_rows,
        .ne12 = (int32_t)n_tokens,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = src1_row_bytes,
        .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes,
        .ne0 = (int32_t)src0_rows,
        .ne1 = (int32_t)selected_experts,
        .nb1 = (uint64_t)src0_rows * sizeof(float),
        .nr0 = (int32_t)nr0,
    };
}

static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args(
        uint32_t src0_cols,
        uint32_t src0_experts,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens) {
    const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float);
    return (ds4_gpu_mul_mm_id_map_args) {
        .ne02 = (int32_t)src0_experts,
        .ne10 = (int32_t)src0_cols,
        .ne11 = (int32_t)src1_expert_rows,
        .nb11 = src1_row_bytes,
        .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes,
        .ne21 = (int32_t)n_tokens,
        .ne20 = (int32_t)selected_experts,
        .nb21 = (uint64_t)selected_experts * sizeof(int32_t),
    };
}

static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args(
        uint32_t src0_cols,
        uint32_t src0_rows,
        uint32_t src0_experts,
        uint64_t src0_row_bytes,
        uint64_t src0_expert_bytes,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens) {
    return ds4_gpu_make_mul_mm_id_args_src1_size(src0_cols,
                                                   src0_rows,
                                                   src0_experts,
                                                   src0_row_bytes,
                                                   src0_expert_bytes,
                                                   src1_expert_rows,
                                                   selected_experts,
                                                   n_tokens,
                                                   sizeof(float));
}

static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size(
        uint32_t src0_cols,
        uint32_t src0_rows,
        uint32_t src0_experts,
        uint64_t src0_row_bytes,
        uint64_t src0_expert_bytes,
        uint32_t src1_expert_rows,
        uint32_t selected_experts,
        uint32_t n_tokens,
        uint32_t src1_elem_size) {
    const uint64_t src1_row_bytes = (uint64_t)src0_cols * src1_elem_size;
    return (ds4_gpu_mul_mm_id_args) {
        .ne00 = (int32_t)src0_cols,
        .ne02 = (int32_t)src0_experts,
        .nb01 = src0_row_bytes,
        .nb02 = src0_expert_bytes,
        .nb03 = (uint64_t)src0_experts * src0_expert_bytes,
        .ne11 = (int32_t)src1_expert_rows,
        .nb10 = src1_elem_size,
        .nb11 = src1_row_bytes,
        .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes,
        .nb13 = (uint64_t)n_tokens * (uint64_t)src1_expert_rows * src1_row_bytes,
        .ne20 = (int32_t)selected_experts,
        .ne21 = (int32_t)n_tokens,
        .ne0 = (int32_t)src0_rows,
        .ne1 = (int32_t)selected_experts,
        .r2 = 1,
        .r3 = 1,
    };
}

static uint32_t ds4_gpu_routed_mv_nr0(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_Q8_0:    return 2;
    case DS4_METAL_TENSOR_Q8_K:    return 2;
    case DS4_METAL_TENSOR_Q4_K:    return 2;
    case DS4_METAL_TENSOR_Q2_K:
    case DS4_METAL_TENSOR_IQ2_XXS: return 4;
    default:                       return 0;
    }
}

static const char *ds4_gpu_metal_tensor_type_name(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_IQ2_XXS: return "iq2_xxs";
    case DS4_METAL_TENSOR_Q2_K:    return "q2_k";
    case DS4_METAL_TENSOR_Q4_K:    return "q4_k";
    case DS4_METAL_TENSOR_Q5_K:    return "q5_k";
    case DS4_METAL_TENSOR_Q6_K:    return "q6_k";
    default:                       return "unknown";
    }
}

static const char *ds4_gpu_trim_env_value(const char *env, size_t *len_out) {
    if (len_out) *len_out = 0;
    if (!env) return NULL;

    while (isspace((unsigned char)*env)) env++;
    size_t n = strlen(env);
    while (n > 0 && isspace((unsigned char)env[n - 1])) n--;
    if (len_out) *len_out = n;
    return env;
}

static bool ds4_gpu_profile_layer_value_match(const char *env, uint32_t layer_index) {
    size_t env_len = 0;
    env = ds4_gpu_trim_env_value(env, &env_len);
    if (!env || env_len == 0) return true;
    if (ds4_gpu_env_value_eq(env, env_len, "all")) return true;

    const char *p = env;
    const char *end_env = env + env_len;
    while (p < end_env) {
        while (p < end_env && (*p == ' ' || *p == '\t' || *p == ',')) p++;
        if (p >= end_env) break;

        char *end = NULL;
        const unsigned long first = strtoul(p, &end, 10);
        if (end == p || end > end_env || first > UINT32_MAX) return false;

        unsigned long last = first;
        p = end;
        if (p < end_env && *p == '-') {
            p++;
            last = strtoul(p, &end, 10);
            if (end == p || end > end_env || last > UINT32_MAX) return false;
            p = end;
        }

        if (first <= layer_index && layer_index <= last) return true;
        while (p < end_env && (*p == ' ' || *p == '\t')) p++;
        if (p < end_env && *p != ',') return false;
    }
    return false;
}

static bool ds4_gpu_stage_profile_enabled_for_layer(const char *flag_env_name,
                                                    const char *layer_env_name,
                                                    uint32_t    layer_index) {
    size_t flag_len = 0;
    const char *flag = ds4_gpu_trim_env_value(getenv(flag_env_name), &flag_len);
    if (!flag) return false;

    size_t layer_len = 0;
    const char *layer = ds4_gpu_trim_env_value(getenv(layer_env_name), &layer_len);
    const bool has_layer_filter = layer && layer_len != 0;

    if (flag_len != 0) {
        if (ds4_gpu_env_value_eq(flag, flag_len, "0") ||
            ds4_gpu_env_value_eq(flag, flag_len, "false") ||
            ds4_gpu_env_value_eq(flag, flag_len, "no") ||
            ds4_gpu_env_value_eq(flag, flag_len, "off")) {
            return false;
        }
        if (!has_layer_filter &&
            !ds4_gpu_env_value_eq(flag, flag_len, "1") &&
            !ds4_gpu_env_value_eq(flag, flag_len, "true") &&
            !ds4_gpu_env_value_eq(flag, flag_len, "yes") &&
            !ds4_gpu_env_value_eq(flag, flag_len, "on") &&
            !ds4_gpu_env_value_eq(flag, flag_len, "all")) {
            return ds4_gpu_profile_layer_value_match(flag, layer_index);
        }
    }

    return ds4_gpu_profile_layer_value_match(layer, layer_index);
}

static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) {
    if (type == DS4_METAL_TENSOR_Q8_0) {
        return 32u * 2u * sizeof(float);
    }
    if (type == DS4_METAL_TENSOR_IQ2_XXS) {
        return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t);
    }
    return 0;
}

static NSUInteger ds4_gpu_routed_mv_nsg(uint32_t type) {
    return type == DS4_METAL_TENSOR_Q8_0 ? 4u : 2u;
}

static bool ds4_gpu_routed_mv_rows_per_group_is_nr0(uint32_t type) {
    return type == DS4_METAL_TENSOR_Q8_0;
}

static id<MTLComputePipelineState> ds4_gpu_routed_mv_pipeline(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_Q8_0:
        return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4);
    case DS4_METAL_TENSOR_Q8_K:
        return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_K_f32", 2);
    case DS4_METAL_TENSOR_IQ2_XXS: return g_moe_mul_mv_id_iq2_xxs_pipeline;
    case DS4_METAL_TENSOR_Q2_K:    return g_moe_mul_mv_id_q2_k_pipeline;
    case DS4_METAL_TENSOR_Q4_K:    return g_moe_mul_mv_id_q4_k_pipeline;
    default:                       return nil;
    }
}

static id<MTLComputePipelineState> ds4_gpu_routed_mm_pipeline(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_Q8_0:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false);
    case DS4_METAL_TENSOR_Q8_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f32", false);
    case DS4_METAL_TENSOR_IQ2_XXS:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32", false);
    case DS4_METAL_TENSOR_Q2_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f32", false);
    case DS4_METAL_TENSOR_Q4_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false);
    case DS4_METAL_TENSOR_Q5_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f32", false);
    case DS4_METAL_TENSOR_Q6_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f32", false);
    default:
        return nil;
    }
}

static id<MTLComputePipelineState> ds4_gpu_routed_mm_addr_pipeline(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_Q2_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f32", false);
    case DS4_METAL_TENSOR_Q4_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f32", false);
    default:
        return nil;
    }
}

static id<MTLComputePipelineState> ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_Q8_0:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f16", false);
    case DS4_METAL_TENSOR_Q8_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f16", false);
    case DS4_METAL_TENSOR_IQ2_XXS:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f16", false);
    case DS4_METAL_TENSOR_Q2_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f16", false);
    case DS4_METAL_TENSOR_Q4_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16", false);
    case DS4_METAL_TENSOR_Q5_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f16", false);
    case DS4_METAL_TENSOR_Q6_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f16", false);
    default:
        return nil;
    }
}

static id<MTLComputePipelineState> ds4_gpu_routed_mm_addr_f16_rhs_pipeline(uint32_t type) {
    switch (type) {
    case DS4_METAL_TENSOR_Q2_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f16", false);
    case DS4_METAL_TENSOR_Q4_K:
        return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f16", false);
    default:
        return nil;
    }
}

static int ds4_gpu_encode_mul_mv_id(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:src0 offset:src0_off atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:dst  offset:dst_off  atIndex:3];
    [enc setBuffer:ids  offset:ids_off  atIndex:4];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_attn_out_low_q8_direct(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !src0 || !src1 || !dst ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
        return 0;
    }

    /* Two row conventions in the classic matvec family: Q8_0 k-splits one
     * nr0-row group across all nsg simdgroups (cross-simdgroup reduce), so a
     * threadgroup covers nr0 rows; Q4_K gives each simdgroup its own nr0
     * rows, covering nr0*nsg.  Dispatching Q8 with the Q4 stride leaves
     * (nsg-1)/nsg of the output rows unwritten. */
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ?
        (NSUInteger)args->nr0 : (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:src0 offset:src0_off atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:dst  offset:dst_off  atIndex:3];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_id_pair(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        id<MTLBuffer>               src0_a,
        NSUInteger                  src0_a_off,
        id<MTLBuffer>               src0_b,
        NSUInteger                  src0_b_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:src0_a offset:src0_a_off atIndex:1];
    [enc setBuffer:src0_b offset:src0_b_off atIndex:2];
    [enc setBuffer:src1   offset:src1_off   atIndex:3];
    [enc setBuffer:dst_a  offset:dst_a_off  atIndex:4];
    [enc setBuffer:dst_b  offset:dst_b_off  atIndex:5];
    [enc setBuffer:ids    offset:ids_off    atIndex:6];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_id_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        id<MTLBuffer>               src0_a,
        NSUInteger                  src0_a_off,
        id<MTLBuffer>               src0_b,
        NSUInteger                  src0_b_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act ||
        !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    [enc setBuffer:src0_a  offset:src0_a_off  atIndex:2];
    [enc setBuffer:src0_b  offset:src0_b_off  atIndex:3];
    [enc setBuffer:src1    offset:src1_off    atIndex:4];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:5];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:6];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
    [enc setBuffer:ids     offset:ids_off     atIndex:8];
    [enc setBuffer:weights offset:weights_off atIndex:9];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        DS4MetalQ4ExpertTable      *gate_table,
        DS4MetalQ4ExpertTable      *up_table,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0,
        bool                        queue_residency) {
    if (!cb || !pipeline || !args || !act || !gate_table || !up_table ||
        !gate_table.argumentBuffer || !up_table.argumentBuffer ||
        !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    [enc setBuffer:gate_table.argumentBuffer offset:0 atIndex:2];
    [enc setBuffer:up_table.argumentBuffer   offset:0 atIndex:3];
    [enc setBuffer:src1    offset:src1_off    atIndex:4];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:5];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:6];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
    [enc setBuffer:ids     offset:ids_off     atIndex:8];
    [enc setBuffer:weights offset:weights_off atIndex:9];
    if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) ||
        !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) {
        ds4_gpu_end_compute_encoder(cb, enc);
        return 0;
    }
    ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, queue_residency);
    ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, queue_residency);
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        DS4MetalQ4ExpertTable      *gate_table,
        DS4MetalQ4ExpertTable      *up_table,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act || !gate_table || !up_table ||
        !gate_table.addressBuffer || !up_table.addressBuffer ||
        !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    [enc setBuffer:gate_table.addressBuffer offset:0 atIndex:2];
    [enc setBuffer:up_table.addressBuffer   offset:0 atIndex:3];
    [enc setBuffer:src1    offset:src1_off    atIndex:4];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:5];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:6];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
    [enc setBuffer:ids     offset:ids_off     atIndex:8];
    [enc setBuffer:weights offset:weights_off atIndex:9];
    if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) ||
        !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) {
        ds4_gpu_end_compute_encoder(cb, enc);
        return 0;
    }
    if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) {
        ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, false);
        ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, false);
    }
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_table_q4_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        DS4MetalQ4ExpertTable      *table,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        queue_residency) {
    if (!cb || !pipeline || !args || !table || !table.argumentBuffer ||
        !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:table.argumentBuffer offset:0 atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:dst  offset:dst_off  atIndex:3];
    [enc setBuffer:ids  offset:ids_off  atIndex:4];
    if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) {
        ds4_gpu_end_compute_encoder(cb, enc);
        return 0;
    }
    ds4_gpu_use_q4_expert_table_resources(cb, enc, table, queue_residency);
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_q4_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        DS4MetalQ4ExpertTable      *table,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !table || !table.addressBuffer ||
        !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:table.addressBuffer offset:0 atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:dst  offset:dst_off  atIndex:3];
    [enc setBuffer:ids  offset:ids_off  atIndex:4];
    if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) {
        ds4_gpu_end_compute_encoder(cb, enc);
        return 0;
    }
    if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) {
        ds4_gpu_use_q4_expert_table_resources(cb, enc, table, false);
    }
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static uint32_t ds4_gpu_q4_expert_group_size(uint32_t n_total_expert) {
    uint32_t group_size = 32;
    const char *env = getenv("DS4_METAL_Q4_EXPERT_GROUP_SIZE");
    if (env && env[0]) {
        char *end = NULL;
        unsigned long v = strtoul(env, &end, 10);
        if (end != env && *end == '\0' && v > 0 && v <= UINT32_MAX) {
            group_size = (uint32_t)v;
        }
    }
    if (group_size == 0) group_size = 1;
    if (group_size > n_total_expert) group_size = n_total_expert;
    return group_size;
}

static int ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        const ds4_gpu_moe_expert_group_args *group,
        id<MTLBuffer>               src0_a,
        NSUInteger                  src0_a_off,
        id<MTLBuffer>               src0_b,
        NSUInteger                  src0_b_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act || !group ||
        !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0 ||
        group->expert_count == 0) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args  length:sizeof(*args)  atIndex:0];
    [enc setBytes:act   length:sizeof(*act)   atIndex:1];
    [enc setBytes:group length:sizeof(*group) atIndex:2];
    [enc setBuffer:src0_a  offset:src0_a_off  atIndex:3];
    [enc setBuffer:src0_b  offset:src0_b_off  atIndex:4];
    [enc setBuffer:src1    offset:src1_off    atIndex:5];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:6];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:7];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8];
    [enc setBuffer:ids     offset:ids_off     atIndex:9];
    [enc setBuffer:weights offset:weights_off atIndex:10];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group_q4_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_moe_expert_group_args *group,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !group || !src0 || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        group->expert_count == 0) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args  length:sizeof(*args)  atIndex:0];
    [enc setBytes:group length:sizeof(*group) atIndex:1];
    [enc setBuffer:src0 offset:src0_off atIndex:2];
    [enc setBuffer:src1 offset:src1_off atIndex:3];
    [enc setBuffer:dst  offset:dst_off  atIndex:4];
    [enc setBuffer:ids  offset:ids_off  atIndex:5];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_id_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               add_in,
        NSUInteger                  add_in_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 ||
        args->nei0 <= 0 || args->nei0 > 8 || args->nei1 <= 0) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:src0 offset:src0_off atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:dst  offset:dst_off  atIndex:3];
    [enc setBuffer:ids  offset:ids_off  atIndex:4];
    [enc setBuffer:(add_in ? add_in : dst) offset:(add_in ? add_in_off : dst_off) atIndex:5];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_q4_gather_slots6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_q4_gather_slots6_args *args,
        __unsafe_unretained id<MTLBuffer> src_groups[6],
        const NSUInteger            src_group_offsets[6],
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off) {
    if (!cb || !pipeline || !args || !src_groups || !src_group_offsets || !ids || !dst ||
        args->expert_bytes == 0 || (args->expert_bytes & 15u) != 0 ||
        args->group_size == 0 || args->n_slots == 0 || args->n_slots > 6) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if (!src_groups[i]) return 0;
    }

    const uint64_t chunks = args->expert_bytes >> 4;
    if (chunks == 0 || chunks > NSUIntegerMax) return 0;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src_groups[i] offset:src_group_offsets[i] atIndex:1 + i];
    }
    [enc setBuffer:ids offset:ids_off atIndex:7];
    [enc setBuffer:dst offset:dst_off atIndex:8];
    const NSUInteger threads = 256u;
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)chunks + threads - 1u) / threads,
                                          (NSUInteger)args->n_slots,
                                          1)
         threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_slots6_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        __unsafe_unretained id<MTLBuffer> src0_a[6],
        const NSUInteger            src0_a_off[6],
        __unsafe_unretained id<MTLBuffer> src0_b[6],
        const NSUInteger            src0_b_off[6],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off ||
        !src1 || !dst_a || !dst_b || !dst_mid || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if (!src0_a[i] || !src0_b[i]) return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i];
    }
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i];
    }
    [enc setBuffer:src1    offset:src1_off    atIndex:14];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:15];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:16];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17];
    [enc setBuffer:weights offset:weights_off atIndex:18];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_slots6_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        __unsafe_unretained id<MTLBuffer> src0[6],
        const NSUInteger            src0_off[6],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if (!src0[i]) return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
    }
    [enc setBuffer:src1 offset:src1_off atIndex:7];
    [enc setBuffer:dst  offset:dst_off  atIndex:8];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group6_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        __unsafe_unretained id<MTLBuffer> src0_a[6],
        const NSUInteger            src0_a_off[6],
        __unsafe_unretained id<MTLBuffer> src0_b[6],
        const NSUInteger            src0_b_off[6],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off ||
        !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 != 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if (!src0_a[i] || !src0_b[i]) return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i];
    }
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i];
    }
    [enc setBuffer:src1    offset:src1_off    atIndex:14];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:15];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:16];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17];
    [enc setBuffer:ids     offset:ids_off     atIndex:18];
    [enc setBuffer:weights offset:weights_off atIndex:19];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group6_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        __unsafe_unretained id<MTLBuffer> src0[6],
        const NSUInteger            src0_off[6],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 != 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if (!src0[i]) return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    for (uint32_t i = 0; i < 6; i++) {
        [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
    }
    [enc setBuffer:src1 offset:src1_off atIndex:7];
    [enc setBuffer:dst  offset:dst_off  atIndex:8];
    [enc setBuffer:ids  offset:ids_off  atIndex:9];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        ds4_gpu_stream_expert_cache_entry * const *entries,
        uint32_t                    n_entries,
        id<MTLBuffer>               gate_addrs,
        id<MTLBuffer>               up_addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0,
        id<MTLBuffer>               overflow_gate,
        id<MTLBuffer>               overflow_up) {
    if (!cb || !pipeline || !args || !act || !entries ||
        (n_entries == 0 && !overflow_gate) ||
        !gate_addrs || !up_addrs ||
        !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 ||
        args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
        args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }
    for (uint32_t i = 0; i < n_entries; i++) {
        if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0;
    }
    if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
                                                           n_entries,
                                                           0)) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    [enc setBuffer:gate_addrs offset:0 atIndex:2];
    [enc setBuffer:up_addrs   offset:0 atIndex:3];
    [enc setBuffer:src1       offset:src1_off    atIndex:4];
    [enc setBuffer:dst_a      offset:dst_a_off   atIndex:5];
    [enc setBuffer:dst_b      offset:dst_b_off   atIndex:6];
    [enc setBuffer:dst_mid    offset:dst_mid_off atIndex:7];
    [enc setBuffer:ids        offset:ids_off     atIndex:8];
    [enc setBuffer:weights    offset:weights_off atIndex:9];
    for (uint32_t i = 0; i < n_entries; i++) {
        [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead];
        [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead];
    }
    /* Overflow experts are addressed straight into the mapped model views
     * when a layer's unique selected set exceeds the cache budget. */
    if (overflow_gate) [enc useResource:overflow_gate usage:MTLResourceUsageRead];
    if (overflow_up) [enc useResource:overflow_up usage:MTLResourceUsageRead];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_iq2(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        ds4_gpu_stream_expert_cache_entry * const *entries,
        uint32_t                    n_entries,
        id<MTLBuffer>               addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !entries || n_entries == 0 ||
        n_entries > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
        !addrs || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 ||
        args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
        args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }
    for (uint32_t i = 0; i < n_entries; i++) {
        if (!entries[i] || !entries[i]->down_buffer) return 0;
    }
    if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
                                                           n_entries,
                                                           0)) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:addrs offset:0 atIndex:1];
    [enc setBuffer:src1  offset:src1_off atIndex:2];
    [enc setBuffer:dst   offset:dst_off  atIndex:3];
    [enc setBuffer:ids   offset:ids_off  atIndex:4];
    for (uint32_t i = 0; i < n_entries; i++) {
        [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead];
    }
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_q2_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        ds4_gpu_stream_expert_cache_entry * const *entries,
        uint32_t                    n_entries,
        id<MTLBuffer>               addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        id<MTLBuffer>               overflow_down) {
    if (!cb || !pipeline || !args || !entries ||
        (n_entries == 0 && !overflow_down) ||
        !addrs || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }
    for (uint32_t i = 0; i < n_entries; i++) {
        if (!entries[i] || !entries[i]->down_buffer) return 0;
    }
    if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
                                                           n_entries,
                                                           0)) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:addrs offset:0 atIndex:1];
    [enc setBuffer:src1  offset:src1_off atIndex:2];
    [enc setBuffer:dst   offset:dst_off  atIndex:3];
    [enc setBuffer:ids   offset:ids_off  atIndex:4];
    for (uint32_t i = 0; i < n_entries; i++) {
        [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead];
    }
    if (overflow_down) [enc useResource:overflow_down usage:MTLResourceUsageRead];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        const ds4_gpu_stream_expert_split_args *split,
        ds4_gpu_stream_expert_cache_entry * const entries[6],
        id<MTLBuffer>               gate_addrs,
        id<MTLBuffer>               up_addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act || !split || !entries ||
        !gate_addrs || !up_addrs || !src1 || !dst_a || !dst_b || !dst_mid ||
        !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if ((split->active_mask & (1u << i)) == 0) continue;
        if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0;
    }
    if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
                                                           6,
                                                           split->active_mask)) {
        return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args  length:sizeof(*args)  atIndex:0];
    [enc setBytes:act   length:sizeof(*act)   atIndex:1];
    [enc setBytes:split length:sizeof(*split) atIndex:2];
    [enc setBuffer:gate_addrs offset:0 atIndex:3];
    [enc setBuffer:up_addrs   offset:0 atIndex:4];
    [enc setBuffer:src1       offset:src1_off    atIndex:5];
    [enc setBuffer:dst_a      offset:dst_a_off   atIndex:6];
    [enc setBuffer:dst_b      offset:dst_b_off   atIndex:7];
    [enc setBuffer:dst_mid    offset:dst_mid_off atIndex:8];
    [enc setBuffer:ids        offset:ids_off     atIndex:9];
    [enc setBuffer:weights    offset:weights_off atIndex:10];
    for (uint32_t i = 0; i < 6; i++) {
        if ((split->active_mask & (1u << i)) == 0) continue;
        [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead];
        [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead];
    }
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_stream_expert_split_args *split,
        ds4_gpu_stream_expert_cache_entry * const entries[6],
        id<MTLBuffer>               addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !split || !entries || !addrs ||
        !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 <= 0 || args->ne02 > 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 6; i++) {
        if ((split->active_mask & (1u << i)) == 0) continue;
        if (!entries[i] || !entries[i]->down_buffer) return 0;
    }
    if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
                                                           6,
                                                           split->active_mask)) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args  length:sizeof(*args)  atIndex:0];
    [enc setBytes:split length:sizeof(*split) atIndex:1];
    [enc setBuffer:addrs offset:0 atIndex:2];
    [enc setBuffer:src1  offset:src1_off atIndex:3];
    [enc setBuffer:dst   offset:dst_off  atIndex:4];
    [enc setBuffer:ids   offset:ids_off  atIndex:5];
    for (uint32_t i = 0; i < 6; i++) {
        if ((split->active_mask & (1u << i)) == 0) continue;
        [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead];
    }
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group8_pair_swiglu(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
        __unsafe_unretained id<MTLBuffer> src0_a[8],
        const NSUInteger            src0_a_off[8],
        __unsafe_unretained id<MTLBuffer> src0_b[8],
        const NSUInteger            src0_b_off[8],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst_a,
        NSUInteger                  dst_a_off,
        id<MTLBuffer>               dst_b,
        NSUInteger                  dst_b_off,
        id<MTLBuffer>               dst_mid,
        NSUInteger                  dst_mid_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off ||
        !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 != 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 8; i++) {
        if (!src0_a[i] || !src0_b[i]) return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBytes:act  length:sizeof(*act)  atIndex:1];
    for (uint32_t i = 0; i < 8; i++) {
        [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i];
    }
    for (uint32_t i = 0; i < 8; i++) {
        [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:10 + i];
    }
    [enc setBuffer:src1    offset:src1_off    atIndex:18];
    [enc setBuffer:dst_a   offset:dst_a_off   atIndex:19];
    [enc setBuffer:dst_b   offset:dst_b_off   atIndex:20];
    [enc setBuffer:dst_mid offset:dst_mid_off atIndex:21];
    [enc setBuffer:ids     offset:ids_off     atIndex:22];
    [enc setBuffer:weights offset:weights_off atIndex:23];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group8_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        __unsafe_unretained id<MTLBuffer> src0[8],
        const NSUInteger            src0_off[8],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 != 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 8; i++) {
        if (!src0[i]) return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    for (uint32_t i = 0; i < 8; i++) {
        [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
    }
    [enc setBuffer:src1 offset:src1_off atIndex:9];
    [enc setBuffer:dst  offset:dst_off  atIndex:10];
    [enc setBuffer:ids  offset:ids_off  atIndex:11];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group24_id(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        __unsafe_unretained id<MTLBuffer> src0[24],
        const NSUInteger            src0_off[24],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg,
        bool                        rows_per_group_is_nr0) {
    if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 != 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 24; i++) {
        if (!src0[i]) return 0;
    }

    const NSUInteger nr0 = (NSUInteger)args->nr0;
    const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
    const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    for (uint32_t i = 0; i < 24; i++) {
        [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
    }
    [enc setBuffer:src1 offset:src1_off atIndex:25];
    [enc setBuffer:dst  offset:dst_off  atIndex:26];
    [enc setBuffer:ids  offset:ids_off  atIndex:27];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mv_group24_sum6(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mv_id_args *args,
        __unsafe_unretained id<MTLBuffer> src0[24],
        const NSUInteger            src0_off[24],
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
        args->ne02 != 384) {
        return 0;
    }
    for (uint32_t i = 0; i < 24; i++) {
        if (!src0[i]) return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
    const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    for (uint32_t i = 0; i < 24; i++) {
        [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
    }
    [enc setBuffer:src1 offset:src1_off atIndex:25];
    [enc setBuffer:dst  offset:dst_off  atIndex:26];
    [enc setBuffer:ids  offset:ids_off  atIndex:27];
    if (threadgroup_bytes != 0) {
        [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    }
    [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
         threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mm_id(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> map_pipeline,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_map_args *map_args,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off) {
    if (!cb || !map_pipeline || !mm_pipeline || !map_args || !mm_args ||
        !src0 || !src1 || !dst || !ids ||
        mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
        mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
        return 0;
    }

    return ds4_gpu_encode_mul_mm_id_map(cb,
                                          map_pipeline,
                                          map_args,
                                          mm_args,
                                          ids,
                                          ids_off) &&
           ds4_gpu_encode_mul_mm_id_mapped(cb,
                                             mm_pipeline,
                                             mm_args,
                                             src0,
                                             src0_off,
                                             src1,
                                             src1_off,
                                             dst,
                                             dst_off);
}

static int ds4_gpu_encode_mul_mm_id_map(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> map_pipeline,
        const ds4_gpu_mul_mm_id_map_args *map_args,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               ids,
        NSUInteger                  ids_off) {
    if (!cb || !map_pipeline || !map_args || !mm_args || !ids ||
        mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
        return 0;
    }

    const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
    const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
    if (tpe_bytes > NSUIntegerMax - hids_bytes) return 0;
    if (!ds4_gpu_ensure_scratch_buffer(&g_moe_id_map_buffer,
                                         &g_moe_id_map_bytes,
                                         tpe_bytes + hids_bytes,
                                         "ds4_moe_id_map")) {
        return 0;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:map_pipeline];
    [enc setBytes:map_args length:sizeof(*map_args) atIndex:0];
    [enc setBuffer:ids offset:ids_off atIndex:1];
    [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2];
    [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:3];
    [enc setThreadgroupMemoryLength:(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne20 * sizeof(uint16_t) atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
         threadsPerThreadgroup:MTLSizeMake((NSUInteger)mm_args->ne02, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mm_id_mapped_tile(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes) {
    if (!cb || !mm_pipeline || !mm_args || !src0 || !src1 || !dst ||
        !g_moe_id_map_buffer ||
        mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
        mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
        return 0;
    }
    /*
     * The routed MoE grouped matmul uses the legacy 32-token expert-major tile.
     * The removed TensorOps variant was not semantically stable on evals, so keep
     * this encoder tied to the tested simdgroup kernel shape.
     */
    const NSUInteger tile_n = 32u;
    const bool use_resource_hints =
        getenv("DS4_METAL_MOE_MM_ID_USE_RESOURCES") != NULL &&
        getenv("DS4_METAL_DISABLE_MOE_MM_ID_USE_RESOURCES") == NULL;

    const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
    const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
    if (tpe_bytes > NSUIntegerMax - hids_bytes ||
        g_moe_id_map_bytes < tpe_bytes + hids_bytes) {
        return 0;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:mm_pipeline];
    [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
    [enc setBuffer:src0 offset:src0_off atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3];
    [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4];
    [enc setBuffer:dst offset:dst_off atIndex:5];
    if (use_resource_hints) {
        [enc useResource:src0 usage:MTLResourceUsageRead];
    }
    [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n,
                                          ((NSUInteger)mm_args->ne0 + 63u) / 64u,
                                          (NSUInteger)mm_args->ne02)
         threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0_addrs,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off,
        NSUInteger                  threadgroup_bytes,
        ds4_gpu_stream_expert_cache_entry * const *resources,
        uint32_t                    resource_count,
        uint32_t                    resource_kind,
        id<MTLBuffer>               overflow_resource) {
    if (!cb || !mm_pipeline || !mm_args || !src0_addrs || !src1 || !dst ||
        !g_moe_id_map_buffer ||
        mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
        mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
        return 0;
    }

    const NSUInteger tile_n = 32u;
    const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
    const NSUInteger hids_bytes =
        (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
    if (tpe_bytes > NSUIntegerMax - hids_bytes ||
        g_moe_id_map_bytes < tpe_bytes + hids_bytes) {
        return 0;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:mm_pipeline];
    [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
    [enc setBuffer:src0_addrs offset:0 atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3];
    [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4];
    [enc setBuffer:dst offset:dst_off atIndex:5];
    [enc useResource:src0_addrs usage:MTLResourceUsageRead];
    for (uint32_t i = 0; resources && i < resource_count; i++) {
        ds4_gpu_stream_expert_cache_entry *entry = resources[i];
        if (!entry) continue;
        id<MTLBuffer> b =
            resource_kind == 0 ? entry->gate_buffer :
            resource_kind == 1 ? entry->up_buffer :
                                 entry->down_buffer;
        if (b) [enc useResource:b usage:MTLResourceUsageRead];
    }
    if (overflow_resource) {
        [enc useResource:overflow_resource usage:MTLResourceUsageRead];
    }
    [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n,
                                          ((NSUInteger)mm_args->ne0 + 63u) / 64u,
                                          (NSUInteger)mm_args->ne02)
         threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        const ds4_gpu_dsv4_moe_swiglu_weight_args *act_args,
        id<MTLBuffer>               gate_src0,
        NSUInteger                  gate_src0_off,
        id<MTLBuffer>               up_src0,
        NSUInteger                  up_src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               mid,
        NSUInteger                  mid_off,
        id<MTLBuffer>               weights,
        NSUInteger                  weights_off) {
    if (!cb || !pipeline || !mm_args || !act_args ||
        !gate_src0 || !up_src0 || !src1 || !mid || !weights ||
        !g_moe_id_map_buffer ||
        mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
        mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
        return 0;
    }

    const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
    const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
    if (tpe_bytes > NSUIntegerMax - hids_bytes ||
        g_moe_id_map_bytes < tpe_bytes + hids_bytes) {
        return 0;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
    [enc setBytes:act_args length:sizeof(*act_args) atIndex:1];
    [enc setBuffer:gate_src0 offset:gate_src0_off atIndex:2];
    [enc setBuffer:up_src0 offset:up_src0_off atIndex:3];
    [enc setBuffer:src1 offset:src1_off atIndex:4];
    [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:5];
    [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:6];
    [enc setBuffer:mid offset:mid_off atIndex:7];
    [enc setBuffer:weights offset:weights_off atIndex:8];
    [enc setThreadgroupMemoryLength:16384u atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + 31u) / 32u,
                                          ((NSUInteger)mm_args->ne0 + 63u) / 64u,
                                          (NSUInteger)mm_args->ne02)
         threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_mul_mm_id_mapped(
        id<MTLCommandBuffer>        cb,
        id<MTLComputePipelineState> mm_pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>               src0,
        NSUInteger                  src0_off,
        id<MTLBuffer>               src1,
        NSUInteger                  src1_off,
        id<MTLBuffer>               dst,
        NSUInteger                  dst_off) {
    return ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                  mm_pipeline,
                                                  mm_args,
                                                  src0,
                                                  src0_off,
                                                  src1,
                                                  src1_off,
                                                  dst,
                                                  dst_off,
                                                  8192u);
}

static int ds4_gpu_encode_attn_out_low_q8_mpp(
        id<MTLCommandBuffer>           cb,
        id<MTLComputePipelineState>    pipeline,
        const ds4_gpu_mul_mm_id_args *mm_args,
        id<MTLBuffer>                  src0,
        NSUInteger                     src0_off,
        id<MTLBuffer>                  src1,
        NSUInteger                     src1_off,
        id<MTLBuffer>                  dst,
        NSUInteger                     dst_off) {
    if (!cb || !pipeline || !mm_args || !src0 || !src1 || !dst ||
        mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
        mm_args->ne02 <= 0 || mm_args->ne1 <= 0 || mm_args->ne21 <= 0) {
        return 0;
    }

    const uint32_t tile_n = DS4_METAL_ATTN_OUT_MPP_TILE_N;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
    [enc setBuffer:src0 offset:src0_off atIndex:1];
    [enc setBuffer:src1 offset:src1_off atIndex:2];
    [enc setBuffer:dst offset:dst_off atIndex:3];
    [enc setThreadgroupMemoryLength:8192u atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + (NSUInteger)tile_n - 1u) / (NSUInteger)tile_n,
                                          ((NSUInteger)mm_args->ne0 + 63u) / 64u,
                                          (NSUInteger)mm_args->ne02)
         threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_swiglu_flat(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        gate,
        NSUInteger           gate_off,
        id<MTLBuffer>        up,
        NSUInteger           up_off,
        id<MTLBuffer>        out,
        NSUInteger           out_off,
        uint32_t             n) {
    if (!cb || !gate || !up || !out || n == 0) return 0;

    ds4_gpu_glu_args args = {
        .ne00 = (int32_t)n,
        .nb01 = (uint64_t)n * sizeof(float),
        .ne10 = (int32_t)n,
        .nb11 = (uint64_t)n * sizeof(float),
        .ne0 = (int32_t)n,
        .nb1 = (uint64_t)n * sizeof(float),
        .i00 = 0,
        .i10 = 0,
        .alpha = 1.0f,
        .limit = 0.0f,
    };
    NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > 256u) nth = 256u;
    if (nth > (NSUInteger)n) nth = (NSUInteger)n;
    if (nth == 0u) nth = 1u;
    const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_swiglu_flat_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:gate offset:gate_off atIndex:1];
    [enc setBuffer:up   offset:up_off   atIndex:2];
    [enc setBuffer:out  offset:out_off  atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_moe_swiglu_weight(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        gate,
        NSUInteger           gate_off,
        id<MTLBuffer>        up,
        NSUInteger           up_off,
        id<MTLBuffer>        mid,
        NSUInteger           mid_off,
        id<MTLBuffer>        weights,
        NSUInteger           weights_off,
        uint32_t             width,
        uint32_t             rows,
        float                clamp_value,
        bool                 mid_f16) {
    if (!cb || !gate || !up || !mid || !weights || width == 0 || rows == 0) return 0;

    id<MTLComputePipelineState> pipeline =
        ds4_gpu_get_pipeline(mid_f16 ? "kernel_dsv4_moe_swiglu_weight_f16" :
                                         "kernel_dsv4_moe_swiglu_weight");
    if (!pipeline) return 0;

    ds4_gpu_dsv4_moe_swiglu_weight_args args = {
        .width = width,
        .rows = rows,
        .gate_row_stride = (uint64_t)width * sizeof(float),
        .up_row_stride = (uint64_t)width * sizeof(float),
        .mid_row_stride = (uint64_t)width * (mid_f16 ? sizeof(uint16_t) : sizeof(float)),
        .weight_stride = sizeof(float),
        .write_clamped = getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ? 1u : 0u,
        .clamp_value = clamp_value,
    };

    NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > 256u) nth = 256u;
    if (nth > width) nth = width;
    if (nth == 0) nth = 1u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:gate    offset:gate_off    atIndex:1];
    [enc setBuffer:up      offset:up_off      atIndex:2];
    [enc setBuffer:mid     offset:mid_off     atIndex:3];
    [enc setBuffer:weights offset:weights_off atIndex:4];
    [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_moe_sum6(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        experts,
        NSUInteger           experts_off,
        id<MTLBuffer>        out,
        NSUInteger           out_off,
        uint32_t             out_dim,
        uint32_t             n_tokens) {
    if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0;

    if (!g_moe_sum6_pipeline) return 0;

    const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float);
    ds4_gpu_dsv4_moe_sum6_args args = {
        .width = out_dim,
        .tokens = n_tokens,
        .src_token_stride = 6u * out_row_bytes,
        .dst_token_stride = out_row_bytes,
    };

    NSUInteger nth = g_moe_sum6_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > 256u) nth = 256u;
    if (nth > out_dim) nth = out_dim;
    if (nth == 0) nth = 1u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_moe_sum6_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:experts offset:experts_off atIndex:1];
    [enc setBuffer:out     offset:out_off     atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_moe_sum8(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        experts,
        NSUInteger           experts_off,
        id<MTLBuffer>        out,
        NSUInteger           out_off,
        uint32_t             out_dim,
        uint32_t             n_tokens) {
    if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0;

    if (!g_moe_sum8_pipeline) return 0;

    const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float);
    ds4_gpu_dsv4_moe_sum6_args args = {
        .width = out_dim,
        .tokens = n_tokens,
        .src_token_stride = 8u * out_row_bytes,
        .dst_token_stride = out_row_bytes,
    };

    NSUInteger nth = g_moe_sum8_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > 256u) nth = 256u;
    if (nth > out_dim) nth = out_dim;
    if (nth == 0) nth = 1u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_moe_sum8_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:experts offset:experts_off atIndex:1];
    [enc setBuffer:out     offset:out_off     atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static ds4_gpu_bin_args ds4_gpu_make_moe_add_args(
        uint32_t out_dim,
        uint32_t n_tokens,
        uint64_t src0_token_stride,
        uint64_t src1_token_stride,
        uint64_t dst_token_stride) {
    return (ds4_gpu_bin_args) {
        .ne00 = (int32_t)out_dim,
        .ne01 = (int32_t)n_tokens,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = src0_token_stride,
        .nb02 = (uint64_t)n_tokens * src0_token_stride,
        .nb03 = (uint64_t)n_tokens * src0_token_stride,
        .ne10 = (int32_t)out_dim,
        .ne11 = (int32_t)n_tokens,
        .ne12 = 1,
        .ne13 = 1,
        .nb10 = sizeof(float),
        .nb11 = src1_token_stride,
        .nb12 = (uint64_t)n_tokens * src1_token_stride,
        .nb13 = (uint64_t)n_tokens * src1_token_stride,
        .ne0 = (int32_t)out_dim,
        .ne1 = (int32_t)n_tokens,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = dst_token_stride,
        .nb2 = (uint64_t)n_tokens * dst_token_stride,
        .nb3 = (uint64_t)n_tokens * dst_token_stride,
        .offs = 0,
        .o1 = { 0 },
    };
}

static int ds4_gpu_encode_moe_sum_experts(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        experts,
        NSUInteger           experts_off,
        id<MTLBuffer>        out,
        NSUInteger           out_off,
        uint32_t             out_dim,
        uint32_t             n_expert,
        uint32_t             n_tokens) {
    if (!cb || !experts || !out || out_dim == 0 || n_expert < 2 || n_tokens == 0) return 0;

    const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float);
    const uint64_t expert_token_stride = (uint64_t)n_expert * out_row_bytes;

    if (n_expert == 6 &&
        ds4_gpu_encode_moe_sum6(cb,
                                  experts,
                                  experts_off,
                                  out,
                                  out_off,
                                  out_dim,
                                  n_tokens)) {
        return 1;
    }

    if (n_expert == 8 &&
        ds4_gpu_encode_moe_sum8(cb,
                                  experts,
                                  experts_off,
                                  out,
                                  out_off,
                                  out_dim,
                                  n_tokens)) {
        return 1;
    }

    ds4_gpu_bin_args first =
        ds4_gpu_make_moe_add_args(out_dim, n_tokens, expert_token_stride, expert_token_stride, out_row_bytes);
    if (!ds4_gpu_encode_bin_f32_rows(cb,
                                       g_add_pipeline,
                                       &first,
                                       experts,
                                       experts_off,
                                       experts,
                                       experts_off + (NSUInteger)out_row_bytes,
                                       out,
                                       out_off)) {
        return 0;
    }

    ds4_gpu_bin_args accum =
        ds4_gpu_make_moe_add_args(out_dim, n_tokens, out_row_bytes, expert_token_stride, out_row_bytes);
    for (uint32_t slot = 2; slot < n_expert; slot++) {
        if (!ds4_gpu_encode_bin_f32_rows(cb,
                                           g_add_pipeline,
                                           &accum,
                                           out,
                                           out_off,
                                           experts,
                                           experts_off + (NSUInteger)((uint64_t)slot * out_row_bytes),
                                           out,
                                           out_off)) {
            return 0;
        }
    }
    return 1;
}

static int ds4_gpu_encode_get_rows_i32_token_rows(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        table,
        NSUInteger           table_off,
        id<MTLBuffer>        tokens,
        NSUInteger           tokens_off,
        const int32_t       *token_inline,
        id<MTLBuffer>        selected,
        NSUInteger           selected_off,
        uint32_t             hash_rows,
        uint32_t             n_cols,
        uint32_t             n_tokens) {
    if (!cb || !table || !selected || hash_rows == 0 || n_cols == 0 || n_tokens == 0) return 0;
    if (!tokens && !token_inline) return 0;

    const uint64_t table_row_bytes = (uint64_t)n_cols * sizeof(int32_t);
    const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
    ds4_gpu_get_rows_args args = {
        .ne00t = (int64_t)n_cols,
        .ne00 = (int64_t)n_cols,
        .nb01 = table_row_bytes,
        .nb02 = (uint64_t)hash_rows * table_row_bytes,
        .nb03 = (uint64_t)hash_rows * table_row_bytes,
        .ne10 = (int32_t)n_tokens,
        .nb10 = sizeof(int32_t),
        .nb11 = token_bytes,
        .nb12 = token_bytes,
        .nb1 = table_row_bytes,
        .nb2 = (uint64_t)n_tokens * table_row_bytes,
        .nb3 = (uint64_t)n_tokens * table_row_bytes,
    };

    NSUInteger nth = (NSUInteger)n_cols;
    const NSUInteger max_threads = g_get_rows_i32_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth > max_threads) nth = max_threads;
    if (nth == 0) nth = 1u;
    const NSUInteger nw0 = ((NSUInteger)n_cols + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_get_rows_i32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:table offset:table_off atIndex:1];
    if (tokens) {
        [enc setBuffer:tokens offset:tokens_off atIndex:2];
    } else {
        [enc setBytes:token_inline length:sizeof(*token_inline) atIndex:2];
    }
    [enc setBuffer:selected offset:selected_off atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_get_rows_f32_router_weights(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        probs,
        NSUInteger           probs_off,
        id<MTLBuffer>        selected,
        NSUInteger           selected_off,
        id<MTLBuffer>        weights,
        NSUInteger           weights_off,
        uint32_t             n_expert,
        uint32_t             n_expert_used,
        uint32_t             n_tokens) {
    if (!cb || !probs || !selected || !weights || n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0;

    const uint64_t probs_token_bytes = (uint64_t)n_expert * sizeof(float);
    const uint64_t selected_row_bytes = (uint64_t)n_expert_used * sizeof(int32_t);
    const uint64_t weights_row_bytes = (uint64_t)n_expert_used * sizeof(float);
    ds4_gpu_get_rows_args args = {
        .ne00t = 1,
        .ne00 = 1,
        .nb01 = sizeof(float),
        .nb02 = probs_token_bytes,
        .nb03 = (uint64_t)n_tokens * probs_token_bytes,
        .ne10 = (int64_t)n_expert_used,
        .nb10 = sizeof(int32_t),
        .nb11 = selected_row_bytes,
        .nb12 = (uint64_t)n_tokens * selected_row_bytes,
        .nb1 = sizeof(float),
        .nb2 = weights_row_bytes,
        .nb3 = (uint64_t)n_tokens * weights_row_bytes,
    };

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_get_rows_f32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:probs offset:probs_off atIndex:1];
    [enc setBuffer:selected offset:selected_off atIndex:2];
    [enc setBuffer:weights offset:weights_off atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_expert_used, n_tokens, 1)
         threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_sum_rows_f32(
        id<MTLCommandBuffer> cb,
        id<MTLBuffer>        src,
        NSUInteger           src_off,
        id<MTLBuffer>        dst,
        NSUInteger           dst_off,
        uint32_t             width,
        uint32_t             rows) {
    if (!cb || !src || !dst || width == 0 || rows == 0) return 0;

    const uint64_t src_row_bytes = (uint64_t)width * sizeof(float);
    ds4_gpu_kargs_sum_rows args = {
        .ne00 = (int64_t)width,
        .ne01 = (int64_t)rows,
        .ne02 = 1,
        .ne03 = 1,
        .nb00 = sizeof(float),
        .nb01 = src_row_bytes,
        .nb02 = (uint64_t)rows * src_row_bytes,
        .nb03 = (uint64_t)rows * src_row_bytes,
        .ne0 = 1,
        .ne1 = (int64_t)rows,
        .ne2 = 1,
        .ne3 = 1,
        .nb0 = sizeof(float),
        .nb1 = sizeof(float),
        .nb2 = (uint64_t)rows * sizeof(float),
        .nb3 = (uint64_t)rows * sizeof(float),
    };

    NSUInteger nth = 32u;
    const NSUInteger max_threads = g_sum_rows_f32_f32_pipeline.maxTotalThreadsPerThreadgroup;
    while (nth < (NSUInteger)args.ne00 && nth < max_threads) nth *= 2u;
    if (nth > max_threads) nth = max_threads;
    if (nth > (NSUInteger)args.ne00) nth = (NSUInteger)args.ne00;
    if (nth == 0) nth = 1u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_sum_rows_f32_f32_pipeline];
    [enc setBytes:&args length:sizeof(args) atIndex:0];
    [enc setBuffer:src offset:src_off atIndex:1];
    [enc setBuffer:dst offset:dst_off atIndex:2];
    [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
    [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);
    return 1;
}

static int ds4_gpu_encode_router_select(
        id<MTLCommandBuffer>  cb,
        ds4_gpu_tensor     *selected,
        ds4_gpu_tensor     *weights,
        ds4_gpu_tensor     *probs,
        id<MTLBuffer>         logitsbuf,
        NSUInteger            logits_off,
        id<MTLBuffer>         biasbuf,
        NSUInteger            bias_off,
        id<MTLBuffer>         hashbuf,
        NSUInteger            hash_off,
        id<MTLBuffer>         tokensbuf,
        NSUInteger            tokens_off,
        const int32_t        *single_token,
        uint32_t              hash_rows,
        uint32_t              n_tokens,
        uint32_t              n_expert,
        uint32_t              n_expert_used,
        float                 expert_weight_scale,
        bool                  has_bias,
        bool                  hash_mode) {
    id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
    id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
    id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
    const NSUInteger selected_off = ds4_gpu_tensor_offset(selected);
    const NSUInteger weights_off = ds4_gpu_tensor_offset(weights);
    const NSUInteger probs_off = ds4_gpu_tensor_offset(probs);

    if (!cb || !selectedbuf || !weightsbuf || !probsbuf || !logitsbuf ||
        n_tokens == 0 || n_expert == 0 || n_expert_used == 0) return 0;

    const NSUInteger probs_bytes = (NSUInteger)n_tokens * (NSUInteger)n_expert * sizeof(float);
    const bool flash_router_fast_path =
        n_expert == 256u &&
        n_expert_used == 6u &&
        fabsf(expert_weight_scale - 1.5f) <= 1.0e-6f;

    int ok = 0;
    if (flash_router_fast_path &&
        !g_quality_mode && n_tokens == 1 &&
        getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) {
        const bool force_simd_weights_fusion =
            getenv("DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION") != NULL;
        const bool use_simd_finalize =
            !hash_mode &&
            g_dsv4_router_finalize_one_simd_pipeline != nil &&
            g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u &&
            g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u &&
            (ds4_gpu_device_name_contains("M3") ||
             ds4_gpu_device_name_contains("M5") ||
             getenv("DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE") != NULL ||
             force_simd_weights_fusion) &&
            getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE") == NULL;
        const bool use_simd_weights_fusion =
            use_simd_finalize &&
            g_dsv4_router_finalize_weights_one_simd_pipeline != nil &&
            g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u &&
            g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u &&
            (ds4_gpu_device_name_contains("M3") ||
             ds4_gpu_device_name_contains("M5") ||
             force_simd_weights_fusion) &&
            getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION") == NULL;
        id<MTLComputePipelineState> softplus_sqrt_pipeline =
            ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline,
                                    "kernel_dsv4_softplus_sqrt_f32_4");
        id<MTLComputePipelineState> router_finalize_pipeline =
            ds4_gpu_hot_pipeline(
                use_simd_weights_fusion
                    ? g_dsv4_router_finalize_weights_one_simd_pipeline
                    : use_simd_finalize
                    ? g_dsv4_router_finalize_one_simd_pipeline
                    : g_dsv4_router_finalize_one_pipeline,
                use_simd_weights_fusion
                    ? "kernel_dsv4_router_finalize_weights_one_simd"
                    : use_simd_finalize
                    ? "kernel_dsv4_router_finalize_one_simd"
                    : "kernel_dsv4_router_finalize_one");
        id<MTLComputePipelineState> router_weights_pipeline = use_simd_weights_fusion
            ? nil
            : ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline,
                                      "kernel_dsv4_router_weights_one");
        if (!softplus_sqrt_pipeline || !router_finalize_pipeline ||
            (!use_simd_weights_fusion && !router_weights_pipeline)) return 0;

        ok = ds4_gpu_encode_unary_f32_rows(cb,
                                             softplus_sqrt_pipeline,
                                             logitsbuf,
                                             logits_off,
                                             probsbuf,
                                             probs_off,
                                             n_expert,
                                             1,
                                             1,
                                             0.0f,
                                             0.0f);
        if (!ok) return 0;

        const bool use_token_buffer = single_token == NULL;
        ds4_gpu_dsv4_router_select_one_args args = {
            .has_bias = has_bias ? 1u : 0u,
            .hash_mode = hash_mode ? 1u : 0u,
            .use_token_buffer = use_token_buffer ? 1u : 0u,
            .token = single_token ? (uint32_t)*single_token : 0u,
            .hash_rows = hash_rows,
        };

        const float zero_f32 = 0.0f;
        const int32_t zero_i32 = 0;
        if ((has_bias && !biasbuf) ||
            (hash_mode && !hashbuf) ||
            (use_token_buffer && !tokensbuf)) {
            return 0;
        }

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:router_finalize_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:probsbuf offset:probs_off atIndex:1];
        if (has_bias) {
            [enc setBuffer:biasbuf offset:bias_off atIndex:2];
        } else {
            [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2];
        }
        if (hash_mode) {
            [enc setBuffer:hashbuf offset:hash_off atIndex:3];
        } else {
            [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3];
        }
        if (use_token_buffer) {
            [enc setBuffer:tokensbuf offset:tokens_off atIndex:4];
        } else {
            [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4];
        }
        [enc setBuffer:selectedbuf offset:selected_off atIndex:5];
        if (use_simd_weights_fusion) {
            [enc setBuffer:weightsbuf offset:weights_off atIndex:6];
        }
        const NSUInteger router_finalize_scratch_bytes = use_simd_finalize
            ? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t))
            : 256u * sizeof(float) + 256u * sizeof(int32_t);
        [enc setThreadgroupMemoryLength:router_finalize_scratch_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (use_simd_weights_fusion) return 1;

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:router_weights_pipeline];
        [enc setBuffer:probsbuf offset:probs_off atIndex:0];
        [enc setBuffer:selectedbuf offset:selected_off atIndex:1];
        [enc setBuffer:weightsbuf offset:weights_off atIndex:2];
        [enc dispatchThreads:MTLSizeMake(6, 1, 1)
        threadsPerThreadgroup:MTLSizeMake(6, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        return 1;
    }

    if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) {
        id<MTLComputePipelineState> softplus_sqrt_pipeline =
            ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline,
                                    "kernel_dsv4_softplus_sqrt_f32_4");
        ok = softplus_sqrt_pipeline &&
             ds4_gpu_encode_unary_f32_rows(cb,
                                             softplus_sqrt_pipeline,
                                             logitsbuf,
                                             logits_off,
                                             probsbuf,
                                             probs_off,
                                             n_expert,
                                             1,
                                             1,
                                             0.0f,
                                             0.0f);
    } else {
        ok = ds4_gpu_encode_unary_f32_rows(cb,
                                             g_unary_softplus_pipeline,
                                             logitsbuf,
                                             logits_off,
                                             probsbuf,
                                             probs_off,
                                             n_expert,
                                             n_tokens,
                                             1,
                                             0.0f,
                                             0.0f) &&
             ds4_gpu_encode_unary_f32_rows(cb,
                                             g_unary_sqrt_pipeline,
                                             probsbuf,
                                             probs_off,
                                             probsbuf,
                                             probs_off,
                                             n_expert,
                                             n_tokens,
                                             1,
                                             0.0f,
                                             0.0f);
    }
    if (!ok) return 0;

    if (hash_mode) {
        ok = ds4_gpu_encode_get_rows_i32_token_rows(cb,
                                                      hashbuf,
                                                      hash_off,
                                                      tokensbuf,
                                                      tokens_off,
                                                      single_token,
                                                      selectedbuf,
                                                      selected_off,
                                                      hash_rows,
                                                      n_expert_used,
                                                      n_tokens);
    } else {
        ds4_gpu_tensor *score_tensor = probs;
        DS4MetalTensor *selection_view = nil;

        if (has_bias) {
            if (!biasbuf ||
                !ds4_gpu_ensure_scratch_buffer(&g_router_selection_buffer,
                                                 &g_router_selection_bytes,
                                                 probs_bytes,
                                                 "ds4_router_selection")) {
                return 0;
            }

            ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_expert, n_tokens, n_expert);
            ok = ds4_gpu_encode_bin_f32_rows(cb,
                                               g_add_pipeline,
                                               &add_args,
                                               probsbuf,
                                               probs_off,
                                               biasbuf,
                                               bias_off,
                                               g_router_selection_buffer,
                                               0);
            if (!ok) return 0;

            selection_view = [DS4MetalTensor new];
            selection_view.buffer = g_router_selection_buffer;
            selection_view.offset = 0;
            selection_view.bytes = probs_bytes;
            selection_view.owner = 0;
            score_tensor = (__bridge ds4_gpu_tensor *)selection_view;
        }

        ok = ds4_gpu_indexer_topk_tensor(selected, score_tensor, n_expert, n_tokens, n_expert_used) != 0;
    }
    if (!ok) return 0;

    const bool use_batch_weights_fusion =
        flash_router_fast_path && !g_quality_mode && n_tokens > 1u &&
        g_dsv4_router_weights_batch_pipeline != nil &&
        (ds4_gpu_device_name_contains("M3") ||
         getenv("DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION") != NULL) &&
        getenv("DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION") == NULL &&
        getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL;
    if (use_batch_weights_fusion) {
        const float scale = expert_weight_scale;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_dsv4_router_weights_batch_pipeline];
        [enc setBytes:&scale length:sizeof(scale) atIndex:0];
        [enc setBuffer:probsbuf offset:probs_off atIndex:1];
        [enc setBuffer:selectedbuf offset:selected_off atIndex:2];
        [enc setBuffer:weightsbuf offset:weights_off atIndex:3];
        [enc setThreadgroupMemoryLength:40u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(6, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        return 1;
    }

    if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) {
        id<MTLComputePipelineState> router_weights_pipeline =
            ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline,
                                    "kernel_dsv4_router_weights_one");
        if (!router_weights_pipeline) return 0;
        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:router_weights_pipeline];
        [enc setBuffer:probsbuf offset:probs_off atIndex:0];
        [enc setBuffer:selectedbuf offset:selected_off atIndex:1];
        [enc setBuffer:weightsbuf offset:weights_off atIndex:2];
        [enc dispatchThreads:MTLSizeMake(6, 1, 1)
        threadsPerThreadgroup:MTLSizeMake(6, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        return 1;
    }

    const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float);
    if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer,
                                         &g_router_weight_sum_bytes,
                                         sum_bytes,
                                         "ds4_router_weight_sum")) {
        return 0;
    }

    ok = ds4_gpu_encode_get_rows_f32_router_weights(cb,
                                                      probsbuf,
                                                      probs_off,
                                                      selectedbuf,
                                                      selected_off,
                                                      weightsbuf,
                                                      weights_off,
                                                      n_expert,
                                                      n_expert_used,
                                                      n_tokens) &&
         ds4_gpu_encode_sum_rows_f32(cb,
                                       weightsbuf,
                                       weights_off,
                                       g_router_weight_sum_buffer,
                                       0,
                                       n_expert_used,
                                       n_tokens) &&
         ds4_gpu_encode_unary_f32_rows(cb,
                                         g_unary_clamp_pipeline,
                                         g_router_weight_sum_buffer,
                                         0,
                                         g_router_weight_sum_buffer,
                                         0,
                                         1,
                                         n_tokens,
                                         0,
                                         6.103515625e-5f,
                                         ds4_gpu_positive_infinity());
    if (!ok) return 0;

    ds4_gpu_bin_args div_args = ds4_gpu_make_bin_rowwise_scalar_args(n_expert_used, n_tokens);
    const float scale = expert_weight_scale;
    ds4_gpu_bin_args scale_args = ds4_gpu_make_bin_rows_args(n_expert_used, n_tokens, 1);

    ok = ds4_gpu_encode_bin_f32_rows(cb,
                                       g_bin_div_row_pipeline,
                                       &div_args,
                                       weightsbuf,
                                       weights_off,
                                       g_router_weight_sum_buffer,
                                       0,
                                       weightsbuf,
                                       weights_off);
    if (!ok) return 0;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_bin_mul_scalar_pipeline];
    [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0];
    [enc setBuffer:weightsbuf offset:weights_off atIndex:1];
    [enc setBytes:&scale length:sizeof(scale) atIndex:2];
    [enc setBuffer:weightsbuf offset:weights_off atIndex:3];
    [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)scale_args.ne1,
                                          (NSUInteger)scale_args.ne2,
                                          (NSUInteger)scale_args.ne3)
         threadsPerThreadgroup:MTLSizeMake(ds4_gpu_bin_threads(n_expert_used, g_bin_mul_scalar_pipeline), 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

int ds4_gpu_glm_kv_lora_rms_norm_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *kv_raw,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              n_tokens,
        uint32_t              kv_raw_dim,
        uint32_t              kv_lora_dim,
        float                 eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !kv_raw || !model_map ||
        n_tokens == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 ||
        kv_lora_dim > kv_raw_dim || (kv_lora_dim & 3u) != 0 ||
        !isfinite(eps) || eps < 0.0f) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(kv_raw);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float);
        const uint64_t weight_bytes = (uint64_t)kv_lora_dim * sizeof(float);
        if (!rawbuf || !outbuf ||
            ds4_gpu_tensor_bytes(kv_raw) < raw_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal GLM KV RMS norm received undersized buffers\n");
            return 0;
        }
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal GLM KV RMS norm weight range is outside the mapped model\n");
            return 0;
        }

        const bool exact_decode_weight_view =
            n_tokens == 1u &&
            weight_bytes <= (1ull << 20) &&
            getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL;
        uint64_t weight_inner = 0;
        id<MTLBuffer> weightbuf = exact_decode_weight_view ?
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           weight_offset,
                                           weight_bytes,
                                           &weight_inner) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     weight_offset,
                                     weight_bytes,
                                     &weight_inner);
        if (!weightbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_kv_lora_rms_norm_pipeline,
                                 "kernel_glm_kv_lora_rms_norm");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_kv_lora_rms_norm_args args = {
            .n_tokens = n_tokens,
            .kv_raw_dim = kv_raw_dim,
            .kv_lora_dim = kv_lora_dim,
            .eps = eps,
        };
        const NSUInteger nth = ds4_gpu_rms_norm_threads(kv_lora_dim);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV RMS norm")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_k_b_project_typed_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *kv_norm,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              weight_type,
        uint32_t              n_tokens,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              n_head) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !kv_norm || !model_map ||
        n_tokens == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || n_head == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv_norm);
        uint64_t row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM k_b projection received unsupported weight type\n");
            return 0;
        }
        const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
        const uint64_t weight_bytes = weight_rows * row_bytes;
        const uint64_t kv_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float);
        if (!outbuf || !kvbuf ||
            ds4_gpu_tensor_bytes(kv_norm) < kv_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes) {
            fprintf(stderr, "ds4: Metal GLM k_b projection received undersized buffers\n");
            return 0;
        }
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal GLM k_b projection range is outside the mapped model\n");
            return 0;
        }
        const NSUInteger q_blocks = ((NSUInteger)qk_nope + 31u) / 32u;
        if (q_blocks > 8u) {
            fprintf(stderr, "ds4: Metal GLM k_b projection q width is too large for the tiled kernel\n");
            return 0;
        }

        uint64_t weight_inner = 0;
        id<MTLBuffer> weightbuf =
            ds4_gpu_wrap_model_range(model_map,
                                      model_size,
                                      weight_offset,
                                      weight_bytes,
                                      &weight_inner);
        if (!weightbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_k_b_project_pipeline,
                                 "kernel_glm_k_b_project_q8_0");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_k_b_project_args args = {
            .n_tokens = n_tokens,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .n_head = n_head,
            .row_bytes = (uint32_t)row_bytes,
            .weight_type = weight_type,
            .pad1 = 0,
            .pad2 = 0,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
        [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1)
             threadsPerThreadgroup:MTLSizeMake(32, q_blocks, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM k_b projection")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_k_b_project_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *kv_norm,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              n_tokens,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              n_head) {
    return ds4_gpu_glm_k_b_project_typed_tensor(out,
                                                kv_norm,
                                                model_map,
                                                model_size,
                                                weight_offset,
                                                DS4_METAL_TENSOR_Q8_0,
                                                n_tokens,
                                                kv_lora_dim,
                                                qk_nope,
                                                n_head);
}

int ds4_gpu_glm_store_compact_kv_tensor(
        ds4_gpu_tensor       *kv_lora_cache,
        ds4_gpu_tensor       *k_rope_cache,
        const ds4_gpu_tensor *kv_norm,
        const ds4_gpu_tensor *kv_raw,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_cap,
        uint32_t              kv_raw_dim,
        uint32_t              kv_lora_dim,
        uint32_t              qk_rope,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw ||
        n_tokens == 0 || cache_cap == 0 ||
        kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 ||
        kv_lora_dim > kv_raw_dim ||
        qk_rope > kv_raw_dim - kv_lora_dim ||
        pos0 > cache_cap || n_tokens > cache_cap - pos0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        id<MTLBuffer> kvnormbuf = ds4_gpu_tensor_buffer(kv_norm);
        id<MTLBuffer> kvrawbuf = ds4_gpu_tensor_buffer(kv_raw);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t kv_cache_bytes =
            (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes =
            (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        const uint64_t kv_norm_bytes =
            (uint64_t)n_tokens * kv_lora_dim * sizeof(float);
        const uint64_t kv_raw_bytes =
            (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
        if (!kvcachebuf || !ropecachebuf || !kvnormbuf || !kvrawbuf ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
            ds4_gpu_tensor_bytes(kv_norm) < kv_norm_bytes ||
            ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) {
            fprintf(stderr, "ds4: Metal GLM compact KV store received undersized buffers\n");
            return 0;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_store_compact_kv_pipeline,
                                 "kernel_glm_store_compact_kv");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_store_compact_kv_args args = {
            .pos0 = pos0,
            .n_tokens = n_tokens,
            .cache_cap = cache_cap,
            .kv_raw_dim = kv_raw_dim,
            .kv_lora_dim = kv_lora_dim,
            .qk_rope = qk_rope,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .pad1 = 0,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:kvnormbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:1];
        [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:2];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 2, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM compact KV store")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
        ds4_gpu_tensor       *q_out,
        const ds4_gpu_tensor *q,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              q_weight_offset,
        uint32_t              q_n,
        ds4_gpu_tensor       *kv_lora_cache,
        ds4_gpu_tensor       *k_rope_cache,
        const ds4_gpu_tensor *kv_raw,
        uint64_t              kv_weight_offset,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_cap,
        uint32_t              kv_raw_dim,
        uint32_t              kv_lora_dim,
        uint32_t              qk_rope,
        bool                  cache_f16,
        float                 eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw ||
        !model_map || n_tokens == 0 || cache_cap == 0 ||
        q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 ||
        (q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 ||
        kv_lora_dim > kv_raw_dim ||
        qk_rope > kv_raw_dim - kv_lora_dim ||
        pos0 > cache_cap || n_tokens > cache_cap - pos0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> qoutbuf = ds4_gpu_tensor_buffer(q_out);
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        id<MTLBuffer> kvrawbuf = ds4_gpu_tensor_buffer(kv_raw);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float);
        const uint64_t kv_weight_bytes = (uint64_t)kv_lora_dim * sizeof(float);
        const uint64_t kv_cache_bytes =
            (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes =
            (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        const uint64_t kv_raw_bytes =
            (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
        if (!qbuf || !qoutbuf || !kvcachebuf || !ropecachebuf || !kvrawbuf ||
            ds4_gpu_tensor_bytes(q) < q_row_bytes * n_tokens ||
            ds4_gpu_tensor_bytes(q_out) < q_row_bytes * n_tokens ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
            ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) {
            fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store received undersized buffers\n");
            return 0;
        }
        if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset ||
            kv_weight_offset > model_size || kv_weight_bytes > model_size - kv_weight_offset) {
            fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store weight range is outside the mapped model\n");
            return 0;
        }

        uint64_t q_weight_inner = 0;
        id<MTLBuffer> q_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                              q_weight_offset, q_row_bytes,
                                                              &q_weight_inner);
        if (!q_weightbuf) return 0;
        uint64_t kv_weight_inner = 0;
        id<MTLBuffer> kv_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                               kv_weight_offset, kv_weight_bytes,
                                                               &kv_weight_inner);
        if (!kv_weightbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_qkv_norm_store_compact_kv_pipeline,
                                 "kernel_glm_qkv_norm_store_compact_kv");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_qkv_norm_store_compact_kv_args args = {
            .pos0 = pos0,
            .n_tokens = n_tokens,
            .cache_cap = cache_cap,
            .q_n = q_n,
            .q_n4 = q_n / 4u,
            .kv_raw_dim = kv_raw_dim,
            .kv_lora_dim = kv_lora_dim,
            .kv_lora_n4 = kv_lora_dim / 4u,
            .qk_rope = qk_rope,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .eps = eps,
            .pad0 = 0,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:q_weightbuf offset:(NSUInteger)q_weight_inner atIndex:2];
        [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3];
        [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:4];
        [enc setBuffer:kv_weightbuf offset:(NSUInteger)kv_weight_inner atIndex:5];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:6];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:7];
        [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 3, 1)
             threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM fused q/kv norm compact store")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_store_indexer_k_tensor(
        ds4_gpu_tensor       *indexer_key_cache,
        const ds4_gpu_tensor *raw_k,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint64_t              bias_offset,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_cap,
        uint32_t              head_dim,
        uint32_t              rot_dim,
        uint32_t              n_ctx_orig,
        float                 eps,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!indexer_key_cache || !raw_k || !model_map ||
        n_tokens == 0 || cache_cap == 0 ||
        head_dim == 0 || rot_dim == 0 ||
        rot_dim > head_dim || (rot_dim & 1u) != 0 ||
        pos0 > cache_cap || n_tokens > cache_cap - pos0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_k);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t cache_bytes =
            (uint64_t)cache_cap * head_dim * cache_elem_bytes;
        const uint64_t raw_bytes =
            (uint64_t)n_tokens * head_dim * sizeof(float);
        const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
        if (!cachebuf || !rawbuf ||
            ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes ||
            ds4_gpu_tensor_bytes(raw_k) < raw_bytes) {
            fprintf(stderr, "ds4: Metal GLM indexer K store received undersized buffers\n");
            return 0;
        }
        if (weight_offset > model_size || norm_bytes > model_size - weight_offset ||
            bias_offset > model_size || norm_bytes > model_size - bias_offset) {
            fprintf(stderr, "ds4: Metal GLM indexer K norm range is outside the mapped model\n");
            return 0;
        }

        uint64_t weight_inner = 0;
        uint64_t bias_inner = 0;
        id<MTLBuffer> weightbuf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                     weight_offset, norm_bytes,
                                     &weight_inner);
        id<MTLBuffer> biasbuf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                     bias_offset, norm_bytes,
                                     &bias_inner);
        if (!weightbuf || !biasbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_store_indexer_k_pipeline,
                                 "kernel_glm_store_indexer_k");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_store_indexer_k_args args = {
            .pos0 = pos0,
            .n_tokens = n_tokens,
            .cache_cap = cache_cap,
            .head_dim = head_dim,
            .rot_dim = rot_dim,
            .n_ctx_orig = n_ctx_orig,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .pad0 = 0,
            .eps = eps,
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
            .pad1 = 0.0f,
        };
        const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2];
        [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:3];
        [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:4];
        [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer K store")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_build_kv_cache_tensor(
        ds4_gpu_tensor       *key_cache,
        ds4_gpu_tensor       *value_cache,
        const ds4_gpu_tensor *kv_raw,
        const ds4_gpu_tensor *k_nope,
        const ds4_gpu_tensor *value,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_cap,
        uint32_t              n_head,
        uint32_t              kv_raw_dim,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    if (!key_cache || !value_cache || !kv_raw || !k_nope || !value ||
        n_tokens == 0 || cache_cap == 0 || n_head == 0 ||
        kv_raw_dim == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_rope == 0 || value_dim == 0 ||
        kv_lora_dim + qk_rope > kv_raw_dim ||
        qk_dim < qk_nope || (qk_rope & 1u) != 0 ||
        pos0 > cache_cap || n_tokens > cache_cap - pos0 ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
        id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(kv_raw);
        id<MTLBuffer> knbuf = ds4_gpu_tensor_buffer(k_nope);
        id<MTLBuffer> vbuf = ds4_gpu_tensor_buffer(value);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
        const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
        const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
        const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float);
        const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
        if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf ||
            ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
            ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes ||
            ds4_gpu_tensor_bytes(kv_raw) < raw_bytes ||
            ds4_gpu_tensor_bytes(k_nope) < kn_bytes ||
            ds4_gpu_tensor_bytes(value) < value_bytes) {
            fprintf(stderr, "ds4: Metal GLM KV cache builder received undersized buffers\n");
            return 0;
        }

        const bool decode_group4 =
            n_tokens == 1u &&
            n_head >= 4u &&
            ds4_gpu_env_bool("DS4_METAL_DISABLE_GLM_DECODE_KV_GROUP4") <= 0;
        id<MTLComputePipelineState> pipeline =
            decode_group4 ?
            ds4_gpu_hot_pipeline(g_glm_build_kv_cache_decode_group4_pipeline,
                                 "kernel_glm_build_kv_cache_decode_group4") :
            ds4_gpu_hot_pipeline(g_glm_build_kv_cache_pipeline,
                                 "kernel_glm_build_kv_cache");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_build_kv_cache_args args = {
            .pos0 = pos0,
            .n_tokens = n_tokens,
            .cache_cap = cache_cap,
            .n_head = n_head,
            .kv_raw_dim = kv_raw_dim,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = value_dim,
            .n_ctx_orig = n_ctx_orig,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .pad0 = 0u,
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1];
        [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2];
        [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3];
        [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4];
        [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5];
        const NSUInteger group_y = decode_group4 ?
            ((NSUInteger)n_head + 3u) / 4u :
            (NSUInteger)n_head;
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, group_y, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV cache build")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_build_kv_cache_flash_tensor(
        ds4_gpu_tensor       *key_cache,
        ds4_gpu_tensor       *value_cache,
        const ds4_gpu_tensor *kv_raw,
        const ds4_gpu_tensor *k_nope,
        const ds4_gpu_tensor *value,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_cap,
        uint32_t              n_head,
        uint32_t              kv_raw_dim,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    if (!key_cache || !value_cache || !kv_raw || !k_nope || !value ||
        pos0 != 0 || n_tokens == 0 || cache_cap == 0 || n_head == 0 ||
        kv_raw_dim == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_rope == 0 || value_dim == 0 ||
        kv_lora_dim + qk_rope > kv_raw_dim ||
        qk_dim < qk_nope || (qk_rope & 1u) != 0 ||
        n_tokens > cache_cap ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
        id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
        id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(kv_raw);
        id<MTLBuffer> knbuf = ds4_gpu_tensor_buffer(k_nope);
        id<MTLBuffer> vbuf = ds4_gpu_tensor_buffer(value);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
        const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
        const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
        const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float);
        const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
        if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf ||
            ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
            ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes ||
            ds4_gpu_tensor_bytes(kv_raw) < raw_bytes ||
            ds4_gpu_tensor_bytes(k_nope) < kn_bytes ||
            ds4_gpu_tensor_bytes(value) < value_bytes) {
            fprintf(stderr, "ds4: Metal GLM staged KV cache builder received undersized buffers\n");
            return 0;
        }

        const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t);
        const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t);
        const NSUInteger key_f16_offset = 0;
        const NSUInteger key_f16_bytes =
            (NSUInteger)n_tokens * (NSUInteger)n_head * q_row_bytes_f16;
        const NSUInteger value_f16_offset = key_f16_bytes;
        const NSUInteger value_f16_bytes =
            (NSUInteger)n_tokens * (NSUInteger)n_head * v_row_bytes_f16;
        const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes;
        if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                           &g_flash_attn_kv_bytes,
                                           kv_f16_bytes,
                                           "ds4_glm_flash_attn_kv_f16")) {
            return 0;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_build_kv_cache_flash_pipeline,
                                 "kernel_glm_build_kv_cache_flash");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_build_kv_cache_args args = {
            .pos0 = pos0,
            .n_tokens = n_tokens,
            .cache_cap = cache_cap,
            .n_head = n_head,
            .kv_raw_dim = kv_raw_dim,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = value_dim,
            .n_ctx_orig = n_ctx_orig,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .pad0 = 0u,
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1];
        [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2];
        [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3];
        [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4];
        [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5];
        [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:6];
        [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:7];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM staged KV cache build")) return 0;
    }

    return 1;
}

static int ds4_gpu_glm_attention_flash_tensor_impl(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *key_cache,
        const ds4_gpu_tensor *value_cache,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_len,
        uint32_t              cache_cap,
        uint32_t              n_head,
        uint32_t              qk_dim,
        uint32_t              value_dim,
        bool                  cache_f16,
        int                   kv_pre_staged) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !key_cache || !value_cache ||
        n_tokens == 0 || cache_len == 0 || cache_cap == 0 ||
        n_head == 0 || qk_dim != 256u || value_dim != 256u ||
        cache_len > cache_cap ||
        pos0 > cache_len || n_tokens > cache_len - pos0 ||
        cache_len > ds4_gpu_glm_flash_attention_max_cache_len()) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
        id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
        const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
        const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
        if (!headsbuf || !qbuf || !keybuf || !valbuf ||
            ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
            ds4_gpu_tensor_bytes(value_cache) < value_bytes) {
            fprintf(stderr, "ds4: Metal GLM FlashAttention received undersized buffers\n");
            return 0;
        }
        const uint64_t key_elems = (uint64_t)cache_len * n_head * qk_dim;
        const uint64_t value_elems = (uint64_t)cache_len * n_head * value_dim;
        if (key_elems > UINT32_MAX || value_elems > UINT32_MAX) {
            return 0;
        }

        const uint32_t nqptg = 8;
        const uint32_t ncpsg = 64;
        const uint32_t nsg = 4;
        const bool has_kvpad = (cache_len % ncpsg) != 0;
        const bool bc_mask = (n_tokens % nqptg) != 0;
        const NSUInteger q_row_bytes = (NSUInteger)qk_dim * sizeof(float);
        const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t);
        const NSUInteger v_row_bytes = (NSUInteger)value_dim * sizeof(float);
        const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t);
        const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)cache_len * sizeof(uint16_t);
        const NSUInteger key_f16_offset = 0;
        const NSUInteger key_f16_bytes =
            (NSUInteger)cache_len * (NSUInteger)n_head * q_row_bytes_f16;
        const NSUInteger value_f16_offset = key_f16_bytes;
        const NSUInteger value_f16_bytes =
            (NSUInteger)cache_len * (NSUInteger)n_head * v_row_bytes_f16;
        const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes;
        const NSUInteger pad_bytes = has_kvpad
            ? (NSUInteger)ncpsg * ((NSUInteger)n_head * (q_row_bytes_f16 + v_row_bytes_f16) +
                                   (NSUInteger)n_tokens * sizeof(uint16_t))
            : 1u;
        const NSUInteger nblk0 = ((NSUInteger)cache_len + ncpsg - 1u) / ncpsg;
        const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg;
        const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);

        id<MTLBuffer> mask_buffer =
            ds4_gpu_glm_prefill_mask_buffer(pos0, n_tokens, cache_len, mask_bytes);
        if (!mask_buffer) return 0;
        if (kv_pre_staged) {
            if (!g_flash_attn_kv_buffer || g_flash_attn_kv_bytes < kv_f16_bytes) {
                fprintf(stderr, "ds4: GLM staged FlashAttention KV scratch is missing\n");
                return 0;
            }
        } else if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
                                                  &g_flash_attn_kv_bytes,
                                                  kv_f16_bytes,
                                                  "ds4_glm_flash_attn_kv_f16")) {
            return 0;
        }
        if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
                                             &g_flash_attn_pad_bytes,
                                             pad_bytes,
                                             "ds4_glm_flash_attn_pad") ||
            !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
                                             &g_flash_attn_blk_bytes,
                                             blk_bytes,
                                             "ds4_glm_flash_attn_blk")) {
            return 0;
        }

        id<MTLComputePipelineState> pad_pipeline = nil;
        if (has_kvpad) {
            pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
            if (!pad_pipeline) return 0;
        }
        id<MTLComputePipelineState> blk_pipeline =
            ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
        id<MTLComputePipelineState> attn_pipeline =
            ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk256_dv256",
                                              true, false, false, false, has_kvpad, bc_mask,
                                              (int32_t)qk_dim,
                                              (int32_t)value_dim,
                                              (int32_t)nsg);
        if (!blk_pipeline || !attn_pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        if (!kv_pre_staged) {
            const bool copied = cache_f16 ?
                (ds4_gpu_encode_cpy_f16_f16_3d(cb,
                                               keybuf,
                                               ds4_gpu_tensor_offset(key_cache),
                                               g_flash_attn_kv_buffer,
                                               key_f16_offset,
                                               qk_dim,
                                               cache_len,
                                               n_head,
                                               (uint64_t)n_head * q_row_bytes_f16,
                                               q_row_bytes_f16,
                                               q_row_bytes_f16,
                                               (uint64_t)cache_len * q_row_bytes_f16) &&
                 ds4_gpu_encode_cpy_f16_f16_3d(cb,
                                               valbuf,
                                               ds4_gpu_tensor_offset(value_cache),
                                               g_flash_attn_kv_buffer,
                                               value_f16_offset,
                                               value_dim,
                                               cache_len,
                                               n_head,
                                               (uint64_t)n_head * v_row_bytes_f16,
                                               v_row_bytes_f16,
                                               v_row_bytes_f16,
                                               (uint64_t)cache_len * v_row_bytes_f16)) :
                (ds4_gpu_encode_cpy_f32_f16_3d(cb,
                                               keybuf,
                                               ds4_gpu_tensor_offset(key_cache),
                                               g_flash_attn_kv_buffer,
                                               key_f16_offset,
                                               qk_dim,
                                               cache_len,
                                               n_head,
                                               (uint64_t)n_head * q_row_bytes,
                                               q_row_bytes,
                                               q_row_bytes_f16,
                                               (uint64_t)cache_len * q_row_bytes_f16) &&
                 ds4_gpu_encode_cpy_f32_f16_3d(cb,
                                               valbuf,
                                               ds4_gpu_tensor_offset(value_cache),
                                               g_flash_attn_kv_buffer,
                                               value_f16_offset,
                                               value_dim,
                                               cache_len,
                                               n_head,
                                               (uint64_t)n_head * v_row_bytes,
                                               v_row_bytes,
                                               v_row_bytes_f16,
                                               (uint64_t)cache_len * v_row_bytes_f16));
            if (!copied) {
                return 0;
            }
        }

        if (has_kvpad) {
            ds4_gpu_flash_attn_pad_args pad_args = {
                .ne11 = (int32_t)cache_len,
                .ne_12_2 = (int32_t)n_head,
                .ne_12_3 = 1,
                .nb11 = q_row_bytes_f16,
                .nb12 = (uint64_t)cache_len * q_row_bytes_f16,
                .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16,
                .nb21 = v_row_bytes_f16,
                .nb22 = (uint64_t)cache_len * v_row_bytes_f16,
                .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16,
                .ne31 = (int32_t)n_tokens,
                .ne32 = 1,
                .ne33 = 1,
                .nb31 = (uint64_t)cache_len * sizeof(uint16_t),
                .nb32 = mask_bytes,
                .nb33 = mask_bytes,
            };

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:pad_pipeline];
            [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
            [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:1];
            [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:2];
            [enc setBuffer:mask_buffer offset:0 atIndex:3];
            [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
            [enc dispatchThreadgroups:MTLSizeMake(ncpsg, n_head, 1)
                 threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);
        }

        ds4_gpu_flash_attn_blk_args blk_args = {
            .ne01 = (int32_t)n_tokens,
            .ne30 = (int32_t)cache_len,
            .ne31 = (int32_t)n_tokens,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)cache_len * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:blk_pipeline];
        [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
        [enc setBuffer:mask_buffer offset:0 atIndex:1];
        [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        ds4_gpu_flash_attn_vec_args args = {
            .ne01 = (int32_t)n_tokens,
            .ne02 = (int32_t)n_head,
            .ne03 = 1,
            .nb01 = (uint64_t)n_head * q_row_bytes,
            .nb02 = q_row_bytes,
            .nb03 = (uint64_t)n_tokens * n_head * q_row_bytes,
            .ne11 = (int32_t)cache_len,
            .ne_12_2 = (int32_t)n_head,
            .ne_12_3 = 1,
            .ns10 = (int32_t)qk_dim,
            .nb11 = q_row_bytes_f16,
            .nb12 = (uint64_t)cache_len * q_row_bytes_f16,
            .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16,
            .ns20 = (int32_t)value_dim,
            .nb21 = v_row_bytes_f16,
            .nb22 = (uint64_t)cache_len * v_row_bytes_f16,
            .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16,
            .ne31 = (int32_t)n_tokens,
            .ne32 = 1,
            .ne33 = 1,
            .nb31 = (uint64_t)cache_len * sizeof(uint16_t),
            .nb32 = mask_bytes,
            .nb33 = mask_bytes,
            .ne1 = (int32_t)n_head,
            .ne2 = (int32_t)n_tokens,
            .ne3 = 1,
            .scale = 1.0f / sqrtf((float)qk_dim),
            .max_bias = 0.0f,
            .m0 = 0.0f,
            .m1 = 0.0f,
            .n_head_log2 = 0,
            .logit_softcap = 0.0f,
        };

        const NSUInteger padded_v = ds4_gpu_align_up_ns(value_dim, 64u);
        const NSUInteger shared_elems = (NSUInteger)nqptg *
            ((NSUInteger)qk_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
        const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:attn_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:2];
        [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:3];
        [enc setBuffer:mask_buffer offset:0 atIndex:4];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:5];
        [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
        [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
        [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
             threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM FlashAttention")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_attention_flash_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *key_cache,
        const ds4_gpu_tensor *value_cache,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_len,
        uint32_t              cache_cap,
        uint32_t              n_head,
        uint32_t              qk_dim,
        uint32_t              value_dim,
        bool                  cache_f16) {
    return ds4_gpu_glm_attention_flash_tensor_impl(heads,
                                                   q,
                                                   key_cache,
                                                   value_cache,
                                                   pos0,
                                                   n_tokens,
                                                   cache_len,
                                                   cache_cap,
                                                   n_head,
                                                   qk_dim,
                                                   value_dim,
                                                   cache_f16,
                                                   0);
}

int ds4_gpu_glm_attention_flash_staged_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *key_cache,
        const ds4_gpu_tensor *value_cache,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_len,
        uint32_t              cache_cap,
        uint32_t              n_head,
        uint32_t              qk_dim,
        uint32_t              value_dim,
        bool                  cache_f16) {
    if (pos0 != 0 || n_tokens != cache_len) return 0;
    return ds4_gpu_glm_attention_flash_tensor_impl(heads,
                                                   q,
                                                   key_cache,
                                                   value_cache,
                                                   pos0,
                                                   n_tokens,
                                                   cache_len,
                                                   cache_cap,
                                                   n_head,
                                                   qk_dim,
                                                   value_dim,
                                                   cache_f16,
                                                   1);
}

int ds4_gpu_glm_attention_full_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *key_cache,
        const ds4_gpu_tensor *value_cache,
        uint32_t              pos0,
        uint32_t              n_tokens,
        uint32_t              cache_len,
        uint32_t              cache_cap,
        uint32_t              n_head,
        uint32_t              qk_dim,
        uint32_t              value_dim,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !key_cache || !value_cache ||
        n_tokens == 0 || cache_len == 0 || cache_cap == 0 ||
        n_head == 0 || qk_dim == 0 || value_dim == 0 ||
        (qk_dim & 3u) != 0 ||
        cache_len > cache_cap ||
        pos0 > cache_len || n_tokens > cache_len - pos0 ||
        cache_len > ds4_gpu_glm_full_attention_max_cache_len()) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
        id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
        const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
        const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
        if (!headsbuf || !qbuf || !keybuf || !valbuf ||
            ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
            ds4_gpu_tensor_bytes(value_cache) < value_bytes) {
            fprintf(stderr, "ds4: Metal GLM attention received undersized buffers\n");
            return 0;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_attention_full_pipeline,
                                 "kernel_glm_attention_full");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const uint32_t full_attn_mode = 2u;
        ds4_gpu_glm_attention_full_args args = {
            .pos0 = pos0,
            .n_tokens = n_tokens,
            .cache_len = cache_len,
            .cache_cap = cache_cap,
            .n_head = n_head,
            .qk_dim = qk_dim,
            .value_dim = value_dim,
            .pad0 = full_attn_mode,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .pad1 = 0u,
            .pad2 = 0u,
            .scale = 1.0f / sqrtf((float)qk_dim),
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:2];
        [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:3];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4];
        [enc setThreadgroupMemoryLength:(256u + (NSUInteger)cache_len) * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM full attention")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_fill_selected_range_tensor(
        ds4_gpu_tensor *selected,
        uint32_t        n_selected) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || n_selected == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t);
        if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal GLM selected range received undersized buffer\n");
            return 0;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_fill_selected_range_pipeline,
                                 "kernel_glm_fill_selected_range");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_fill_selected_range_args args = {
            .n_selected = n_selected,
        };
        const NSUInteger nth = 256u;
        const NSUInteger n_groups = ((NSUInteger)n_selected + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1];
        [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_fill_selected_range_batch_tensor(
        ds4_gpu_tensor *selected,
        uint32_t        n_tokens,
        uint32_t        pos0,
        uint32_t        n_selected,
        uint32_t        pad_row) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || n_tokens == 0 || n_selected == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        const uint64_t total = (uint64_t)n_tokens * n_selected;
        if (n_tokens != 0 && total / n_tokens != n_selected) return 0;
        if (total > UINT64_MAX / sizeof(uint32_t)) return 0;
        const uint64_t selected_bytes = total * sizeof(uint32_t);
        if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal GLM selected range batch received undersized buffer\n");
            return 0;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_fill_selected_range_batch_pipeline,
                                 "kernel_glm_fill_selected_range_batch");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_fill_selected_range_batch_args args = {
            .n_tokens = n_tokens,
            .pos0 = pos0,
            .n_selected = n_selected,
            .pad_row = pad_row,
        };
        const NSUInteger nth = 256u;
        const NSUInteger n_groups = ((NSUInteger)total + nth - 1u) / nth;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1];
        [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range batch")) return 0;
    }

    return 1;
}

static int ds4_gpu_glm_rope_tail_offset_tensor(
        ds4_gpu_tensor *x,
        uint32_t        n_tokens,
        uint32_t        n_head,
        uint32_t        head_dim,
        uint32_t        rot_dim,
        uint32_t        rot_offset,
        uint32_t        pos0,
        uint32_t        n_ctx_orig,
        float           freq_base,
        float           freq_scale,
        float           ext_factor,
        float           attn_factor,
        float           beta_fast,
        float           beta_slow,
        const char     *label) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 ||
        rot_dim == 0 || rot_offset > head_dim || rot_dim > head_dim - rot_offset ||
        (rot_dim & 1u) != 0 ||
        pos0 > UINT32_MAX - n_tokens ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        const uint64_t bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
        if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
            fprintf(stderr, "ds4: Metal %s received undersized buffer\n",
                    label ? label : "GLM RoPE");
            return 0;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_indexer_rope_tail_pipeline,
                                 "kernel_glm_indexer_rope_tail_f32");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_rope_tail_args args = {
            .n_tokens = n_tokens,
            .n_head = n_head,
            .head_dim = head_dim,
            .rot_dim = rot_dim,
            .rot_offset = rot_offset,
            .pos0 = pos0,
            .n_ctx_orig = n_ctx_orig,
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
        };
        const NSUInteger nth = ds4_gpu_rms_norm_threads(rot_dim);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, label ? label : "GLM RoPE")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_rope_tail_tensor(
        ds4_gpu_tensor *x,
        uint32_t        n_tokens,
        uint32_t        n_head,
        uint32_t        head_dim,
        uint32_t        rot_dim,
        uint32_t        pos0,
        uint32_t        n_ctx_orig,
        float           freq_base,
        float           freq_scale,
        float           ext_factor,
        float           attn_factor,
        float           beta_fast,
        float           beta_slow) {
    if (rot_dim > head_dim) return 0;
    return ds4_gpu_glm_rope_tail_offset_tensor(x,
                                               n_tokens,
                                               n_head,
                                               head_dim,
                                               rot_dim,
                                               head_dim - rot_dim,
                                               pos0,
                                               n_ctx_orig,
                                               freq_base,
                                               freq_scale,
                                               ext_factor,
                                               attn_factor,
                                               beta_fast,
                                               beta_slow,
                                               "GLM RoPE");
}

int ds4_gpu_glm_indexer_rope_tail_tensor(
        ds4_gpu_tensor *x,
        uint32_t        n_tokens,
        uint32_t        n_head,
        uint32_t        head_dim,
        uint32_t        rot_dim,
        uint32_t        pos0,
        uint32_t        n_ctx_orig,
        float           freq_base,
        float           freq_scale,
        float           ext_factor,
        float           attn_factor,
        float           beta_fast,
        float           beta_slow) {
    return ds4_gpu_glm_rope_tail_offset_tensor(x,
                                               n_tokens,
                                               n_head,
                                               head_dim,
                                               rot_dim,
                                               0,
                                               pos0,
                                               n_ctx_orig,
                                               freq_base,
                                               freq_scale,
                                               ext_factor,
                                               attn_factor,
                                               beta_fast,
                                               beta_slow,
                                               "GLM indexer RoPE");
}

int ds4_gpu_glm_indexer_score_one_tensor(
        ds4_gpu_tensor       *scores,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *weights,
        const ds4_gpu_tensor *indexer_key_cache,
        uint32_t              n_rows,
        uint32_t              n_head,
        uint32_t              head_dim,
        float                 scale,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!scores || !q || !weights || !indexer_key_cache ||
        n_rows == 0 || n_head == 0 || head_dim == 0 ||
        !isfinite(scale) || scale <= 0.0f) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> scoresbuf = ds4_gpu_tensor_buffer(scores);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t score_bytes = (uint64_t)n_rows * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
        const uint64_t weights_bytes = (uint64_t)n_head * sizeof(float);
        const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes;
        if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf ||
            ds4_gpu_tensor_bytes(scores) < score_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes ||
            ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) {
            fprintf(stderr, "ds4: Metal GLM indexer score received undersized buffers\n");
            return 0;
        }

        ds4_gpu_glm_indexer_score_one_args args = {
            .n_rows = n_rows,
            .n_head = n_head,
            .head_dim = head_dim,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .scale = scale,
        };

        if (n_head == 32u && head_dim == 128u) {
            id<MTLComputePipelineState> direct_pipeline =
                ds4_gpu_hot_pipeline(g_glm_indexer_score_one_direct_pipeline,
                                     "kernel_glm_indexer_score_one_direct");
            if (!direct_pipeline) return 0;

            int owned = 0;
            id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:direct_pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
            [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
            [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3];
            [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
            [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer direct score")) return 0;
            return 1;
        }

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_indexer_score_one_pipeline,
                                 "kernel_glm_indexer_score_one");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
        [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3];
        [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
        [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer score")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_indexer_scores_batch_tensor(
        ds4_gpu_tensor       *scores,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *weights,
        const ds4_gpu_tensor *indexer_key_cache,
        uint32_t              n_rows,
        uint32_t              n_tokens,
        uint32_t              pos0,
        uint32_t              n_head,
        uint32_t              head_dim,
        float                 scale,
        bool                  cache_f16) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!scores || !q || !weights || !indexer_key_cache ||
        n_rows == 0 || n_tokens == 0 || n_head == 0 || head_dim != 128 ||
        pos0 >= n_rows || n_tokens > n_rows - pos0 ||
        !isfinite(scale) || scale <= 0.0f) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> scoresbuf = ds4_gpu_tensor_buffer(scores);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t score_bytes = (uint64_t)n_rows * n_tokens * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
        const uint64_t weights_bytes = (uint64_t)n_tokens * n_head * sizeof(float);
        const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes;
        if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf ||
            ds4_gpu_tensor_bytes(scores) < score_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes ||
            ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) {
            fprintf(stderr, "ds4: Metal GLM indexer batch scores received undersized buffers\n");
            return 0;
        }

        const bool force_scalar = g_quality_mode;
        const bool use_tiled_f32 = false;
        const bool use_tiled = !force_scalar && n_tokens >= 8u &&
                               n_head == 32u && head_dim == 128u;
        id<MTLComputePipelineState> pipeline =
            use_tiled
                ? ds4_gpu_hot_pipeline(use_tiled_f32 ? g_glm_indexer_scores_tiled_f32_pipeline
                                                     : g_glm_indexer_scores_tiled_pipeline,
                                       use_tiled_f32 ? "kernel_glm_indexer_scores_tiled_f32"
                                                     : "kernel_glm_indexer_scores_tiled")
                : ds4_gpu_hot_pipeline(g_glm_indexer_scores_batch_pipeline,
                                       "kernel_glm_indexer_scores_batch");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_indexer_scores_batch_args args = {
            .n_rows = n_rows,
            .n_tokens = n_tokens,
            .n_head = n_head,
            .head_dim = head_dim,
            .pos0 = pos0,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float),
            .q_head_stride = (uint64_t)head_dim * sizeof(float),
            .weights_token_stride = (uint64_t)n_head * sizeof(float),
            .score_token_stride = (uint64_t)n_rows * sizeof(float),
            .scale = scale,
        };
        const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim);

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
        [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3];
        [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
        if (use_tiled) {
            const NSUInteger q_shared = 8u * 128u;
            const NSUInteger k_shared = 32u * 128u;
            const NSUInteger dot_shared = 8u * 32u;
            if (use_tiled_f32) {
                [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) *
                                                sizeof(float) atIndex:0];
            } else {
                [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) +
                                                dot_shared * sizeof(float) atIndex:0];
            }
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + 31u) / 32u,
                                                  ((NSUInteger)n_tokens + 7u) / 8u,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
        } else {
            [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows,
                                                  (NSUInteger)n_tokens,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer batch scores")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_qk_lowrank_typed_tensor(
        ds4_gpu_tensor       *qk_low,
        const ds4_gpu_tensor *q,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              weight_type,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!qk_low || !q || !model_map ||
        n_head == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_nope > qk_dim) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        uint64_t row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM qk lowrank received unsupported weight type\n");
            return 0;
        }
        const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
        const uint64_t weight_bytes = weight_rows * row_bytes;
        const uint64_t out_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float);
        if (!outbuf || !qbuf ||
            ds4_gpu_tensor_bytes(qk_low) < out_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes) {
            fprintf(stderr, "ds4: Metal GLM qk lowrank received undersized buffers\n");
            return 0;
        }
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal GLM qk lowrank range is outside the mapped model\n");
            return 0;
        }

        uint64_t weight_inner = 0;
        id<MTLBuffer> weightbuf =
            ds4_gpu_wrap_model_range(model_map,
                                      model_size,
                                      weight_offset,
                                      weight_bytes,
                                      &weight_inner);
        if (!weightbuf) return 0;

        const int use_glm52 =
            n_head == 64u &&
            kv_lora_dim == 512u &&
            qk_nope == 192u &&
            qk_dim == 256u &&
            row_bytes == 204u &&
            weight_type == DS4_METAL_TENSOR_Q8_0;
        /* Coalesced simdgroup variant: lanes split the 192-dot so weight
         * reads coalesce, and 2048 threadgroups replace 64. The thread-
         * per-row kernels measured ~7.5x off the weight-bandwidth floor
         * (7.2ms of the decode token by skip-ablation). Covers the GLM 5.2
         * shape for both Q8_0 k_b and the DenseQ4 GGUF's Q4_0 k_b (the Q8
         * fast path above never engaged there). */
        const int use_glm52_sg =
            n_head == 64u &&
            kv_lora_dim == 512u &&
            qk_nope == 192u &&
            qk_dim == 256u &&
            ((weight_type == DS4_METAL_TENSOR_Q8_0 && row_bytes == 204u) ||
             (weight_type == DS4_METAL_TENSOR_Q4_0 && row_bytes == 108u)) &&
            getenv("DS4_METAL_DISABLE_GLM_QKLOW_SG") == NULL &&
            ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline,
                                 "kernel_glm_qk_lowrank_q8_0_glm52_sg") != nil;
        if (getenv("DS4_METAL_GLM_QKLOW_DEBUG")) {
            static int printed = 0;
            if (!printed) {
                printed = 1;
                fprintf(stderr, "ds4: qk_lowrank decode path: use_glm52=%d sg=%d n_head=%u kv=%u nope=%u dim=%u rb=%llu type=%u\n",
                        use_glm52, use_glm52_sg, n_head, kv_lora_dim, qk_nope, qk_dim,
                        (unsigned long long)row_bytes, weight_type);
            }
        }
        id<MTLComputePipelineState> pipeline =
            use_glm52_sg ?
                ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline,
                                     "kernel_glm_qk_lowrank_q8_0_glm52_sg") :
            use_glm52 ?
                ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_pipeline,
                                     "kernel_glm_qk_lowrank_q8_0_glm52") :
                ds4_gpu_hot_pipeline(g_glm_qk_lowrank_pipeline,
                                     "kernel_glm_qk_lowrank_q8_0");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_qk_lowrank_args args = {
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_dim = qk_dim,
            .row_bytes = (uint32_t)row_bytes,
            .weight_type = weight_type,
            .pad1 = 0,
            .pad2 = 0,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3];
        if (use_glm52_sg) {
            /* 8 simdgroups x 2 rows per threadgroup: (64, 512/16) grid. */
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head,
                                                  (NSUInteger)(kv_lora_dim / 16u),
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
        } else {
            if (use_glm52) {
                [enc setThreadgroupMemoryLength:192u * sizeof(float) atIndex:0];
            }
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM qk lowrank")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_qk_lowrank_q8_0_tensor(
        ds4_gpu_tensor       *qk_low,
        const ds4_gpu_tensor *q,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_dim) {
    return ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low,
                                               q,
                                               model_map,
                                               model_size,
                                               weight_offset,
                                               DS4_METAL_TENSOR_Q8_0,
                                               n_head,
                                               kv_lora_dim,
                                               qk_nope,
                                               qk_dim);
}

int ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
        ds4_gpu_tensor       *qk_low,
        const ds4_gpu_tensor *q,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              weight_type,
        uint32_t              n_tokens,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!qk_low || !q || !model_map ||
        n_tokens == 0 || n_head == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_nope > qk_dim) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        uint64_t row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM batch qk lowrank received unsupported weight type\n");
            return 0;
        }
        const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
        const uint64_t weight_bytes = weight_rows * row_bytes;
        const uint64_t out_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        const uint64_t q_bytes =
            (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
        if (!outbuf || !qbuf ||
            ds4_gpu_tensor_bytes(qk_low) < out_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes) {
            fprintf(stderr, "ds4: Metal GLM batch qk lowrank received undersized buffers\n");
            return 0;
        }
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal GLM batch qk lowrank range is outside the mapped model\n");
            return 0;
        }

        uint64_t weight_inner = 0;
        id<MTLBuffer> weightbuf =
            ds4_gpu_wrap_model_range(model_map,
                                      model_size,
                                      weight_offset,
                                      weight_bytes,
                                      &weight_inner);
        if (!weightbuf) return 0;

        const int use_glm52_t4 =
            n_tokens >= 4u &&
            n_head == 64u &&
            kv_lora_dim == 512u &&
            qk_nope == 192u &&
            qk_dim == 256u &&
            row_bytes == 204u &&
            weight_type == DS4_METAL_TENSOR_Q8_0;
        id<MTLComputePipelineState> pipeline =
            use_glm52_t4 ?
                ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline,
                                     "kernel_glm_qk_lowrank_q8_0_batch_glm52_t4") :
                ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline,
                                     "kernel_glm_qk_lowrank_q8_0_batch");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        uint32_t head_base = 0;
        uint32_t head_count = n_head;
        ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count);
        ds4_gpu_glm_qk_lowrank_batch_args args = {
            .n_tokens = n_tokens,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_dim = qk_dim,
            .row_bytes = (uint32_t)row_bytes,
            .weight_type = weight_type,
            .head_base = head_base,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3];
        if (use_glm52_t4) {
            [enc setThreadgroupMemoryLength:4u * 192u * sizeof(float) atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count,
                                                  ((NSUInteger)n_tokens + 3u) / 4u,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        } else {
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count,
                                                  (NSUInteger)n_tokens,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch qk lowrank")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor(
        ds4_gpu_tensor       *qk_low,
        const ds4_gpu_tensor *q,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              n_tokens,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_dim) {
    return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low,
                                                     q,
                                                     model_map,
                                                     model_size,
                                                     weight_offset,
                                                     DS4_METAL_TENSOR_Q8_0,
                                                     n_tokens,
                                                     n_head,
                                                     kv_lora_dim,
                                                     qk_nope,
                                                     qk_dim);
}

int ds4_gpu_glm_value_project_typed_batch_heads_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *lora,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              weight_type,
        uint32_t              n_tokens,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              value_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !lora || !model_map ||
        n_tokens == 0 || n_head == 0 ||
        kv_lora_dim == 0 || value_dim == 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        id<MTLBuffer> lorabuf = ds4_gpu_tensor_buffer(lora);
        uint64_t row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(weight_type, kv_lora_dim, &row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM batch value project received unsupported weight type\n");
            return 0;
        }
        const uint64_t weight_rows = (uint64_t)n_head * value_dim;
        const uint64_t weight_bytes = weight_rows * row_bytes;
        const uint64_t heads_bytes =
            (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
        const uint64_t lora_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        if (!headsbuf || !lorabuf ||
            ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(lora) < lora_bytes) {
            fprintf(stderr, "ds4: Metal GLM batch value project received undersized buffers\n");
            return 0;
        }
        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal GLM batch value project range is outside the mapped model\n");
            return 0;
        }

        uint64_t weight_inner = 0;
        id<MTLBuffer> weightbuf =
            ds4_gpu_wrap_model_range(model_map,
                                      model_size,
                                      weight_offset,
                                      weight_bytes,
                                      &weight_inner);
        if (!weightbuf) return 0;

        const int use_mma =
            n_head == 64u &&
            kv_lora_dim == 512u &&
            value_dim == 256u &&
            row_bytes == 544u &&
            weight_type == DS4_METAL_TENSOR_Q8_0 &&
            n_tokens >= 32u;
        id<MTLComputePipelineState> pipeline =
            use_mma ?
                ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_mma_pipeline,
                                     "kernel_glm_value_project_q8_0_batch_heads_mma") :
                ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_pipeline,
                                     "kernel_glm_value_project_q8_0_batch_heads");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        uint32_t head_base = 0;
        uint32_t head_count = n_head;
        ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count);
        ds4_gpu_glm_qk_lowrank_batch_args args = {
            .n_tokens = n_tokens,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = 0,
            .qk_dim = value_dim,
            .row_bytes = (uint32_t)row_bytes,
            .weight_type = weight_type,
            .head_base = head_base,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
        [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:2];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3];
        if (use_mma) {
            [enc setThreadgroupMemoryLength:16u * 1024u atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tokens + 31u) / 32u,
                                                  ((NSUInteger)value_dim + 63u) / 64u,
                                                  (NSUInteger)head_count)
                 threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
        } else {
            [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float)
                                    atIndex:0];
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count,
                                                  (NSUInteger)n_tokens,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch value project")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *lora,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              weight_offset,
        uint32_t              n_tokens,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              value_dim) {
    return ds4_gpu_glm_value_project_typed_batch_heads_tensor(heads,
                                                              lora,
                                                              model_map,
                                                              model_size,
                                                              weight_offset,
                                                              DS4_METAL_TENSOR_Q8_0,
                                                              n_tokens,
                                                              n_head,
                                                              kv_lora_dim,
                                                              value_dim);
}

int ds4_gpu_glm_attention_indexed_decode_typed_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              value_weight_offset,
        uint32_t              value_weight_type,
        const ds4_gpu_tensor *selected,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
        !model_map || !selected ||
        n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
        n_head == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 ||
        value_dim == 0 || qk_dim < qk_nope ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        uint64_t value_row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM indexed attention received unsupported value type\n");
            return 0;
        }
        const uint64_t value_weight_rows = (uint64_t)n_head * value_dim;
        const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes;
        const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float);
        const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float);
        const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t);
        if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
            ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal GLM indexed attention received undersized buffers\n");
            return 0;
        }
        if (value_weight_offset > model_size ||
            value_weight_bytes > model_size - value_weight_offset) {
            fprintf(stderr, "ds4: Metal GLM indexed attention value range is outside the mapped model\n");
            return 0;
        }

        uint64_t value_inner = 0;
        id<MTLBuffer> valuebuf =
            ds4_gpu_wrap_model_range(model_map,
                                      model_size,
                                      value_weight_offset,
                                      value_weight_bytes,
                                      &value_inner);
        if (!valuebuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_pipeline,
                                 "kernel_glm_attention_indexed_decode");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        ds4_gpu_glm_attention_indexed_decode_args args = {
            .n_selected = n_selected,
            .cache_cap = cache_cap,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = value_dim,
            .n_ctx_orig = n_ctx_orig,
            .value_row_bytes = (uint32_t)value_row_bytes,
            .scale = 1.0f / sqrtf((float)qk_dim),
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
            .value_type = value_weight_type,
        };
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
        [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7];
        const NSUInteger scratch_floats =
            256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim;
        [enc setThreadgroupMemoryLength:scratch_floats * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed attention decode")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_attention_indexed_decode_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              value_weight_offset,
        const ds4_gpu_tensor *selected,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    return ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads,
                                                             q,
                                                             qk_low,
                                                             kv_lora_cache,
                                                             k_rope_cache,
                                                             model_map,
                                                             model_size,
                                                             value_weight_offset,
                                                             DS4_METAL_TENSOR_Q8_0,
                                                             selected,
                                                             n_selected,
                                                             cache_cap,
                                                             cache_f16,
                                                             n_head,
                                                             kv_lora_dim,
                                                             qk_nope,
                                                             qk_rope,
                                                             value_dim,
                                                             n_ctx_orig,
                                                             freq_base,
                                                             freq_scale,
                                                             ext_factor,
                                                             attn_factor,
                                                             beta_fast,
                                                             beta_slow);
}

int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(
        ds4_gpu_tensor       *heads,
        ds4_gpu_tensor       *partial_lora,
        ds4_gpu_tensor       *partial_ms,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              value_weight_offset,
        uint32_t              value_weight_type,
        const ds4_gpu_tensor *selected,
        uint32_t              n_selected,
        bool                  selected_rows_valid,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        uint32_t              block_rows,
        uint32_t              n_blocks,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    const uint32_t needed_blocks =
        block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u;
    if (!heads || !partial_lora || !partial_ms || !q || !qk_low ||
        !kv_lora_cache || !k_rope_cache || !model_map || !selected ||
        n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
        n_head == 0 || (n_head % 8u) != 0 ||
        kv_lora_dim != 512u ||
        qk_nope == 0 || qk_rope != 64u ||
        value_dim == 0 || qk_dim < qk_nope ||
        block_rows == 0u || needed_blocks == 0u ||
        n_blocks < needed_blocks || n_blocks > 64u ||
        !cache_f16 ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        id<MTLBuffer> partial_lorabuf = ds4_gpu_tensor_buffer(partial_lora);
        id<MTLBuffer> partial_msbuf = ds4_gpu_tensor_buffer(partial_ms);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        const uint64_t cache_elem_bytes = sizeof(uint16_t);
        uint64_t value_row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received unsupported value type\n");
            return 0;
        }
        const uint64_t value_weight_rows = (uint64_t)n_head * value_dim;
        const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes;
        const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float);
        const uint64_t partial_lora_bytes =
            (uint64_t)n_blocks * n_head * kv_lora_dim * sizeof(float);
        const uint64_t partial_ms_bytes =
            (uint64_t)n_blocks * n_head * 2u * sizeof(float);
        const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float);
        const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float);
        const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t);
        if (!headsbuf || !partial_lorabuf || !partial_msbuf ||
            !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
            ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(partial_lora) < partial_lora_bytes ||
            ds4_gpu_tensor_bytes(partial_ms) < partial_ms_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received undersized buffers\n");
            return 0;
        }
        if (value_weight_offset > model_size ||
            value_weight_bytes > model_size - value_weight_offset) {
            fprintf(stderr, "ds4: Metal GLM split grouped indexed attention value range is outside the mapped model\n");
            return 0;
        }

        uint64_t value_inner = 0;
        id<MTLBuffer> valuebuf =
            ds4_gpu_wrap_model_range(model_map,
                                      model_size,
                                      value_weight_offset,
                                      value_weight_bytes,
                                      &value_inner);
        if (!valuebuf) return 0;

        const bool use_valid_fullheads =
            selected_rows_valid && (n_head % 8u) == 0u;
        id<MTLComputePipelineState> partial_pipeline =
            ds4_gpu_hot_pipeline(use_valid_fullheads ?
                                     g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline :
                                     g_glm_attention_indexed_decode_split_group8_partial_pipeline,
                                 use_valid_fullheads ?
                                     "kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads" :
                                     "kernel_glm_attention_indexed_decode_split_group8_partial");
        const bool use_reduce16 =
            n_blocks == 16u && block_rows == 128u &&
            n_selected == 2048u && value_dim == 256u;
        id<MTLComputePipelineState> reduce_pipeline =
            ds4_gpu_hot_pipeline(use_reduce16 ?
                                     g_glm_attention_indexed_decode_split_group8_reduce16_pipeline :
                                     g_glm_attention_indexed_decode_split_group8_reduce_pipeline,
                                 use_reduce16 ?
                                     "kernel_glm_attention_indexed_decode_split_group8_reduce16" :
                                     "kernel_glm_attention_indexed_decode_split_group8_reduce");
        if (!partial_pipeline || !reduce_pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_attention_indexed_decode_split_args args = {
            .n_selected = n_selected,
            .cache_cap = cache_cap,
            .cache_f16 = 1u,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = value_dim,
            .n_ctx_orig = n_ctx_orig,
            .value_row_bytes = (uint32_t)value_row_bytes,
            .block_rows = block_rows,
            .n_blocks = n_blocks,
            .scale = 1.0f / sqrtf((float)qk_dim),
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
            .value_type = value_weight_type,
        };
        const NSUInteger stage_rows = 16u;
        const NSUInteger kv_vecs = (NSUInteger)kv_lora_dim / 4u;
        const NSUInteger rope_vecs = (NSUInteger)qk_rope / 4u;
        const NSUInteger partial_scratch_bytes =
            stage_rows * kv_vecs * sizeof(uint16_t) * 4u +
            stage_rows * rope_vecs * sizeof(float) * 4u;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:partial_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
        [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:6];
        [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:7];
        [enc setThreadgroupMemoryLength:partial_scratch_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head / 8u, (NSUInteger)n_blocks, 1)
             threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        const NSUInteger reduce_scratch_floats = 256u + 64u + (NSUInteger)kv_lora_dim;
        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:reduce_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:1];
        [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:2];
        [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:3];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4];
        [enc setThreadgroupMemoryLength:reduce_scratch_floats * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM split grouped indexed attention decode")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor(
        ds4_gpu_tensor       *heads,
        ds4_gpu_tensor       *partial_lora,
        ds4_gpu_tensor       *partial_ms,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              value_weight_offset,
        const ds4_gpu_tensor *selected,
        uint32_t              n_selected,
        bool                  selected_rows_valid,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        uint32_t              block_rows,
        uint32_t              n_blocks,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    return ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(heads,
                                                                          partial_lora,
                                                                          partial_ms,
                                                                          q,
                                                                          qk_low,
                                                                          kv_lora_cache,
                                                                          k_rope_cache,
                                                                          model_map,
                                                                          model_size,
                                                                          value_weight_offset,
                                                                          DS4_METAL_TENSOR_Q8_0,
                                                                          selected,
                                                                          n_selected,
                                                                          selected_rows_valid,
                                                                          cache_cap,
                                                                          cache_f16,
                                                                          n_head,
                                                                          kv_lora_dim,
                                                                          qk_nope,
                                                                          qk_rope,
                                                                          value_dim,
                                                                          n_ctx_orig,
                                                                          block_rows,
                                                                          n_blocks,
                                                                          freq_base,
                                                                          freq_scale,
                                                                          ext_factor,
                                                                          attn_factor,
                                                                          beta_fast,
                                                                          beta_slow);
}

int ds4_gpu_glm_attention_indexed_batch_typed_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              value_weight_offset,
        uint32_t              value_weight_type,
        const ds4_gpu_tensor *selected,
        uint32_t              n_tokens,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
        !model_map || !selected ||
        n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
        n_head == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 ||
        value_dim == 0 || qk_dim < qk_nope ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        uint64_t value_row_bytes = 0;
        if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) {
            fprintf(stderr, "ds4: Metal GLM indexed batch attention received unsupported value type\n");
            return 0;
        }
        const uint64_t value_weight_rows = (uint64_t)n_head * value_dim;
        const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes;
        const uint64_t heads_bytes =
            (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
        const uint64_t q_bytes =
            (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
        const uint64_t low_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        const uint64_t selected_bytes =
            (uint64_t)n_tokens * n_selected * sizeof(uint32_t);
        if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
            ds4_gpu_tensor_bytes(heads) < heads_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal GLM indexed batch attention received undersized buffers\n");
            return 0;
        }
        if (value_weight_offset > model_size ||
            value_weight_bytes > model_size - value_weight_offset) {
            fprintf(stderr, "ds4: Metal GLM indexed batch attention value range is outside the mapped model\n");
            return 0;
        }

        uint64_t value_inner = 0;
        id<MTLBuffer> valuebuf =
            ds4_gpu_wrap_model_range(model_map, model_size,
                                      value_weight_offset,
                                      value_weight_bytes,
                                      &value_inner);
        if (!valuebuf) return 0;

        const uint64_t grouped_attn = (n_tokens >= 128u) ? 8u : 2u;
        const NSUInteger q2_bit_words = ((NSUInteger)cache_cap + 31u) / 32u;
        const NSUInteger q2_scratch_bytes =
            q2_bit_words * sizeof(uint32_t) +
            4u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) +
            8u * (NSUInteger)kv_lora_dim * sizeof(float);
        const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength];
        const bool q2_fits = max_tg_mem == 0 || q2_scratch_bytes <= max_tg_mem;
        const bool use_q2_group4 = grouped_attn == 4u && n_tokens >= 2u && q2_fits;
        const bool use_group8 =
            grouped_attn == 8u ||
            (grouped_attn == 4u && !use_q2_group4);
        const bool use_group2 = grouped_attn == 1u || grouped_attn == 2u;
        id<MTLComputePipelineState> pipeline = use_q2_group4 ?
            ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_q2_group4_pipeline,
                                 "kernel_glm_attention_indexed_batch_q2_group4") :
            (use_group8 ?
            ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline,
                                 "kernel_glm_attention_indexed_batch_group8") :
            (use_group2 ?
            ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group2_pipeline,
                                 "kernel_glm_attention_indexed_batch_group2") :
            ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_pipeline,
                                 "kernel_glm_attention_indexed_batch")));
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_attention_indexed_batch_args args = {
            .n_tokens = n_tokens,
            .n_selected = n_selected,
            .cache_cap = cache_cap,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = value_dim,
            .n_ctx_orig = n_ctx_orig,
            .value_row_bytes = (uint32_t)value_row_bytes,
            .value_type = value_weight_type,
            .scale = 1.0f / sqrtf((float)qk_dim),
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
            .head_base = 0,
        };
        const NSUInteger scratch_bytes = use_q2_group4 ?
            q2_scratch_bytes :
            (use_group8 ?
             (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) +
              8u * (NSUInteger)kv_lora_dim * sizeof(float)) :
            ((use_group2 ?
             (512u + 2u * (NSUInteger)n_selected + 2u * (NSUInteger)kv_lora_dim) :
             (256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim)) * sizeof(float)));

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
        [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6];
        [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7];
        [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
        if (use_q2_group4) {
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 3u) / 4u,
                                                  ((NSUInteger)n_tokens + 1u) / 2u,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
        } else if (use_group8) {
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 7u) / 8u,
                                                  (NSUInteger)n_tokens,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
        } else if (use_group2) {
            [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 1u) / 2u,
                                                  (NSUInteger)n_tokens,
                                                  1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        } else {
            [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1)
                 threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        }
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_attention_indexed_batch_tensor(
        ds4_gpu_tensor       *heads,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const void           *model_map,
        uint64_t              model_size,
        uint64_t              value_weight_offset,
        const ds4_gpu_tensor *selected,
        uint32_t              n_tokens,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              value_dim,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    return ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads,
                                                            q,
                                                            qk_low,
                                                            kv_lora_cache,
                                                            k_rope_cache,
                                                            model_map,
                                                            model_size,
                                                            value_weight_offset,
                                                            DS4_METAL_TENSOR_Q8_0,
                                                            selected,
                                                            n_tokens,
                                                            n_selected,
                                                            cache_cap,
                                                            cache_f16,
                                                            n_head,
                                                            kv_lora_dim,
                                                            qk_nope,
                                                            qk_rope,
                                                            value_dim,
                                                            n_ctx_orig,
                                                            freq_base,
                                                            freq_scale,
                                                            ext_factor,
                                                            attn_factor,
                                                            beta_fast,
                                                            beta_slow);
}

int ds4_gpu_sort_i32_rows_asc_tensor(
        ds4_gpu_tensor       *dst,
        const ds4_gpu_tensor *src,
        uint32_t              row_width,
        uint32_t              n_rows) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!dst || !src || row_width == 0 || n_rows == 0 ||
        (row_width & (row_width - 1u)) != 0) {
        return 0;
    }

    @autoreleasepool {
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline,
                                 "kernel_dsv4_sort_i32_rows_asc");
        if (!pipeline) return 0;

        const uint64_t bytes = (uint64_t)row_width * n_rows * sizeof(int32_t);
        id<MTLBuffer> srcbuf = ds4_gpu_tensor_buffer(src);
        id<MTLBuffer> dstbuf = ds4_gpu_tensor_buffer(dst);
        if (!srcbuf || !dstbuf ||
            ds4_gpu_tensor_bytes(src) < bytes ||
            ds4_gpu_tensor_bytes(dst) < bytes) {
            fprintf(stderr, "ds4: Metal row sort received undersized buffers\n");
            return 0;
        }

        NSUInteger threads = (NSUInteger)row_width;
        const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
        if (max_threads != 0 && threads > max_threads) threads = max_threads;
        if (threads == 0) return 0;

        const NSUInteger scratch_bytes = (NSUInteger)row_width * sizeof(int32_t);
        const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength];
        if (max_tg_mem != 0 && scratch_bytes > max_tg_mem) {
            fprintf(stderr, "ds4: Metal row sort scratch exceeds threadgroup memory limit\n");
            return 0;
        }

        ds4_gpu_dsv4_topk_mask_args args = {
            .ne00 = (int64_t)row_width,
            .ne01 = (int64_t)n_rows,
            .nb00 = sizeof(int32_t),
            .nb01 = (uint64_t)row_width * sizeof(int32_t),
            .ne0 = (int64_t)row_width,
            .ne1 = (int64_t)n_rows,
            .nb0 = sizeof(int32_t),
            .nb1 = (uint64_t)row_width * sizeof(int32_t),
        };

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1];
        [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:2];
        [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "sort i32 rows asc")) return 0;
    }

    return 1;
}

static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(
        ds4_gpu_tensor       *lora_out,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const ds4_gpu_tensor *selected,
        uint32_t              n_tokens,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow,
        bool                  selected_rows_valid) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || !selected ||
        n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
        n_head == 0 || kv_lora_dim == 0 ||
        qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 ||
        qk_dim < qk_nope ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> lorabuf = ds4_gpu_tensor_buffer(lora_out);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
        const uint64_t lora_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        const uint64_t q_bytes =
            (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
        const uint64_t low_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        const uint64_t kv_cache_bytes =
            (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes =
            (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        const uint64_t selected_bytes =
            (uint64_t)n_tokens * n_selected * sizeof(uint32_t);
        if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
            ds4_gpu_tensor_bytes(lora_out) < lora_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes) {
            fprintf(stderr, "ds4: Metal GLM indexed batch attention-lora received undersized buffers\n");
            return 0;
        }

        const bool use_vec_lora =
            cache_f16 && kv_lora_dim == 512u && qk_rope == 64u;
        const bool full_head_groups = (n_head % 8u) == 0u;
        id<MTLComputePipelineState> pipeline = nil;
        if (use_vec_lora && selected_rows_valid && full_head_groups) {
            pipeline = ds4_gpu_hot_pipeline(
                    g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline,
                    "kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads");
        } else if (use_vec_lora && selected_rows_valid) {
            pipeline = ds4_gpu_hot_pipeline(
                    g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline,
                    "kernel_glm_attention_indexed_batch_lora_group8_vec_valid");
        } else if (use_vec_lora) {
            pipeline = ds4_gpu_hot_pipeline(
                    g_glm_attention_indexed_batch_lora_group8_vec_pipeline,
                    "kernel_glm_attention_indexed_batch_lora_group8_vec");
        } else {
            pipeline = ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline,
                                            "kernel_glm_attention_indexed_batch_group8");
        }
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_attention_indexed_batch_args args = {
            .n_tokens = n_tokens,
            .n_selected = n_selected,
            .cache_cap = cache_cap,
            .cache_f16 = cache_f16 ? 1u : 0u,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = kv_lora_dim,
            .n_ctx_orig = n_ctx_orig,
            .value_row_bytes = 0,
            .value_type = 1u,
            .scale = 1.0f / sqrtf((float)qk_dim),
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
            .head_base = 0,
        };
        uint32_t head_count = n_head;
        ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count);
        const NSUInteger scratch_bytes = use_vec_lora ?
            (16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u +
             16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u) :
            (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) +
             8u * (NSUInteger)kv_lora_dim * sizeof(float));

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
        if (use_vec_lora) {
            [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
            [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:6];
        } else {
            [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5];
            [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6];
            [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:7];
        }
        [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u,
                                              (NSUInteger)n_tokens,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention-lora")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
        ds4_gpu_tensor       *lora_out,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        uint32_t              n_tokens,
        uint32_t              pos0,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    const uint32_t qk_dim = qk_nope + qk_rope;
    if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
        n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
        pos0 > n_selected || n_tokens > n_selected - pos0 ||
        n_head == 0 || kv_lora_dim != 512u ||
        qk_nope == 0 || qk_rope != 64u ||
        qk_dim < qk_nope || !cache_f16 ||
        !isfinite(freq_base) || freq_base <= 0.0f ||
        !isfinite(freq_scale) || freq_scale <= 0.0f ||
        !isfinite(ext_factor) || !isfinite(attn_factor) ||
        !isfinite(beta_fast) || !isfinite(beta_slow)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> lorabuf = ds4_gpu_tensor_buffer(lora_out);
        id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
        id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
        id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
        id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
        const uint64_t cache_elem_bytes = sizeof(uint16_t);
        const uint64_t lora_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        const uint64_t q_bytes =
            (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
        const uint64_t low_bytes =
            (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
        const uint64_t kv_cache_bytes =
            (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
        const uint64_t rope_cache_bytes =
            (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
        if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf ||
            ds4_gpu_tensor_bytes(lora_out) < lora_bytes ||
            ds4_gpu_tensor_bytes(q) < q_bytes ||
            ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
            ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
            ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes) {
            fprintf(stderr, "ds4: Metal GLM causal batch attention-lora received undersized buffers\n");
            return 0;
        }

        const bool full_head_groups = (n_head % 8u) == 0u;
        id<MTLComputePipelineState> pipeline = full_head_groups ?
            ds4_gpu_hot_pipeline(
                    g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline,
                    "kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads") :
            ds4_gpu_hot_pipeline(
                    g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline,
                    "kernel_glm_attention_indexed_batch_lora_group8_vec_causal");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_attention_indexed_batch_args args = {
            .n_tokens = n_tokens,
            .n_selected = n_selected,
            .cache_cap = cache_cap,
            .cache_f16 = 1u,
            .n_head = n_head,
            .kv_lora_dim = kv_lora_dim,
            .qk_nope = qk_nope,
            .qk_rope = qk_rope,
            .value_dim = kv_lora_dim,
            .n_ctx_orig = n_ctx_orig,
            .value_row_bytes = 0,
            .value_type = 1u,
            .pos0 = pos0,
            .scale = 1.0f / sqrtf((float)qk_dim),
            .freq_base = freq_base,
            .freq_scale = freq_scale,
            .ext_factor = ext_factor,
            .attn_factor = attn_factor,
            .beta_fast = beta_fast,
            .beta_slow = beta_slow,
            .head_base = 0,
        };
        uint32_t head_count = n_head;
        ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count);
        const NSUInteger scratch_bytes =
            16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u +
            16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
        [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
        [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
        [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
        [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5];
        [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u,
                                              (NSUInteger)n_tokens,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM causal indexed batch attention-lora")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_attention_indexed_batch_lora_tensor(
        ds4_gpu_tensor       *lora_out,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const ds4_gpu_tensor *selected,
        uint32_t              n_tokens,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out,
                                                                 q,
                                                                 qk_low,
                                                                 kv_lora_cache,
                                                                 k_rope_cache,
                                                                 selected,
                                                                 n_tokens,
                                                                 n_selected,
                                                                 cache_cap,
                                                                 cache_f16,
                                                                 n_head,
                                                                 kv_lora_dim,
                                                                 qk_nope,
                                                                 qk_rope,
                                                                 n_ctx_orig,
                                                                 freq_base,
                                                                 freq_scale,
                                                                 ext_factor,
                                                                 attn_factor,
                                                                 beta_fast,
                                                                 beta_slow,
                                                                 false);
}

int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
        ds4_gpu_tensor       *lora_out,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *qk_low,
        const ds4_gpu_tensor *kv_lora_cache,
        const ds4_gpu_tensor *k_rope_cache,
        const ds4_gpu_tensor *selected,
        uint32_t              n_tokens,
        uint32_t              n_selected,
        uint32_t              cache_cap,
        bool                  cache_f16,
        uint32_t              n_head,
        uint32_t              kv_lora_dim,
        uint32_t              qk_nope,
        uint32_t              qk_rope,
        uint32_t              n_ctx_orig,
        float                 freq_base,
        float                 freq_scale,
        float                 ext_factor,
        float                 attn_factor,
        float                 beta_fast,
        float                 beta_slow) {
    return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out,
                                                                 q,
                                                                 qk_low,
                                                                 kv_lora_cache,
                                                                 k_rope_cache,
                                                                 selected,
                                                                 n_tokens,
                                                                 n_selected,
                                                                 cache_cap,
                                                                 cache_f16,
                                                                 n_head,
                                                                 kv_lora_dim,
                                                                 qk_nope,
                                                                 qk_rope,
                                                                 n_ctx_orig,
                                                                 freq_base,
                                                                 freq_scale,
                                                                 ext_factor,
                                                                 attn_factor,
                                                                 beta_fast,
                                                                 beta_slow,
                                                                 true);
}

int ds4_gpu_glm_router_select_tensor(
        ds4_gpu_tensor       *selected,
        ds4_gpu_tensor       *weights,
        ds4_gpu_tensor       *probs,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                bias_offset,
        const ds4_gpu_tensor *logits,
        uint32_t                n_expert,
        uint32_t                n_expert_used,
        float                   expert_weight_scale) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || !weights || !probs || !logits || !model_map ||
        n_expert == 0 || n_expert > 256u ||
        n_expert_used == 0 || n_expert_used > n_expert) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
        if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf ||
            ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) ||
            ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int32_t) ||
            ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) ||
            ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) {
            fprintf(stderr, "ds4: Metal GLM router received undersized buffers\n");
            return 0;
        }

        const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
        if (bias_offset > model_size || bias_bytes > model_size - bias_offset) {
            fprintf(stderr, "ds4: Metal GLM router bias range is outside the mapped model\n");
            return 0;
        }
        const bool exact_bias_view =
            bias_bytes <= (1ull << 20) &&
            getenv("DS4_METAL_DISABLE_DECODE_ROUTER_BIAS_EXACT_VIEWS") == NULL;
        uint64_t bias_inner = 0;
        id<MTLBuffer> biasbuf = exact_bias_view ?
            ds4_gpu_wrap_model_exact_range(model_map,
                                           model_size,
                                           bias_offset,
                                           bias_bytes,
                                           &bias_inner) :
            ds4_gpu_wrap_model_range(model_map,
                                     model_size,
                                     bias_offset,
                                     bias_bytes,
                                     &bias_inner);
        if (!biasbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline,
                                 "kernel_glm_router_select_one");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_router_select_one_args args = {
            .n_expert = n_expert,
            .n_expert_used = n_expert_used,
            .expert_weight_scale = expert_weight_scale,
            .pad0 = 0,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1];
        [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3];
        [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4];
        [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5];
        [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM router select")) return 0;
    }

    return 1;
}

int ds4_gpu_glm_router_select_batch_tensor(
        ds4_gpu_tensor       *selected,
        ds4_gpu_tensor       *weights,
        ds4_gpu_tensor       *probs,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                bias_offset,
        const ds4_gpu_tensor *logits,
        uint32_t                n_expert,
        uint32_t                n_expert_used,
        float                   expert_weight_scale,
        uint32_t                n_tokens) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || !weights || !probs || !logits || !model_map ||
        n_tokens == 0 ||
        n_expert == 0 || n_expert > 256u ||
        n_expert_used == 0 || n_expert_used > n_expert) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
        const uint64_t logits_bytes = (uint64_t)n_tokens * n_expert * sizeof(float);
        const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(int32_t);
        const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(float);
        const uint64_t probs_bytes = (uint64_t)n_tokens * n_expert * sizeof(float);
        if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf ||
            ds4_gpu_tensor_bytes(logits) < logits_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes ||
            ds4_gpu_tensor_bytes(probs) < probs_bytes) {
            fprintf(stderr, "ds4: Metal GLM batch router received undersized buffers\n");
            return 0;
        }

        const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
        if (bias_offset > model_size || bias_bytes > model_size - bias_offset) {
            fprintf(stderr, "ds4: Metal GLM batch router bias range is outside the mapped model\n");
            return 0;
        }
        uint64_t bias_inner = 0;
        id<MTLBuffer> biasbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                         bias_offset, bias_bytes,
                                                         &bias_inner);
        if (!biasbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline,
                                 "kernel_glm_router_select_one");
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        ds4_gpu_glm_router_select_one_args args = {
            .n_expert = n_expert,
            .n_expert_used = n_expert_used,
            .expert_weight_scale = expert_weight_scale,
            .pad0 = 0,
        };

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1];
        [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3];
        [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4];
        [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5];
        [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch router select")) return 0;
    }

    return 1;
}

static bool ds4_gpu_glm_gate_pair_type_supported(
        uint32_t gate_type,
        uint32_t up_type) {
    return gate_type == up_type &&
           (gate_type == DS4_METAL_TENSOR_Q2_K ||
            gate_type == DS4_METAL_TENSOR_Q4_K ||
            gate_type == DS4_METAL_TENSOR_Q5_K);
}

static bool ds4_gpu_glm_down_type_supported(uint32_t down_type) {
    return down_type == DS4_METAL_TENSOR_Q2_K ||
           down_type == DS4_METAL_TENSOR_Q4_K ||
           down_type == DS4_METAL_TENSOR_Q5_K ||
           down_type == DS4_METAL_TENSOR_Q6_K;
}

int ds4_gpu_glm_routed_moe_one_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                up_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                up_expert_bytes,
        uint64_t                up_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        uint32_t                layer_index,
        const ds4_gpu_tensor *x,
        bool                    force_resident) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    /* TP sharding: only the owned contiguous expert range is mapped,
     * so bind from the owned base, validate only its bytes, and tell the
     * kernels the first expert id present at that base. */
    uint32_t first_expert = 0;
    uint32_t n_bind_expert = 0;
    ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
    const int32_t tp_expert_base_host = (int32_t)first_expert;
    gate_offset += (uint64_t)first_expert * gate_expert_bytes;
    up_offset += (uint64_t)first_expert * up_expert_bytes;
    down_offset += (uint64_t)first_expert * down_expert_bytes;

    if (!out || !mid || !model_map || !selected || !weights || !x ||
        n_total_expert == 0 || n_expert == 0 || n_expert > 256u ||
        n_expert > n_total_expert ||
        expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 ||
        gate_expert_bytes == 0 || gate_row_bytes == 0 ||
        up_expert_bytes == 0 || up_row_bytes == 0 ||
        down_expert_bytes == 0 || down_row_bytes == 0 ||
        (expert_in_dim % 256u) != 0 ||
        (expert_mid_dim % 256u) != 0 ||
        !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) ||
        !ds4_gpu_glm_down_type_supported(down_type)) {
        return 0;
    }

    if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes ||
        (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes ||
        (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes ||
        (uint64_t)out_dim > UINT64_MAX / down_row_bytes) {
        fprintf(stderr, "ds4: Metal GLM routed MoE tensor byte size overflow\n");
        return 0;
    }

    const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
    const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
    if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes ||
        up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes ||
        down_expert_bytes != (uint64_t)out_dim * down_row_bytes) {
        fprintf(stderr, "ds4: Metal GLM routed MoE received inconsistent expert strides\n");
        return 0;
    }
    if (layer_index == 3u) {
        ds4_gpu_stream_expert_cache_note_decode_token();
    }
    if (gate_offset > model_size || gate_tensor_bytes > model_size - gate_offset ||
        up_offset > model_size || up_tensor_bytes > model_size - up_offset ||
        down_offset > model_size || down_tensor_bytes > model_size - down_offset) {
        fprintf(stderr, "ds4: Metal GLM routed MoE tensor range is outside the mapped model\n");
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float);
        const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float);
        if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(mid) < mid_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int32_t) ||
            ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) {
            fprintf(stderr, "ds4: Metal GLM routed MoE received undersized activation buffers\n");
            return 0;
        }

        const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K;
        const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K;
        const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K;
        const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K;
        const BOOL down_simd_q4 = down_scalar_q4;
        const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K;
        const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K;
        const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6;
        const BOOL stream_addr_q2 =
            gate_pair_q2 && down_scalar_q2 &&
            g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil &&
            g_glm_q2_k_addr_down_f32_pipeline != nil;
        const BOOL stream_addr_q4 =
            !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 &&
            g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil &&
            g_glm_q4_k_addr_down_f32_pipeline != nil;
        BOOL use_stream_expert_addr_table =
            g_ssd_streaming_mode &&
            !force_resident &&
            getenv("DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE") == NULL &&
            (stream_addr_q2 || stream_addr_q4) &&
            layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
            n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
            n_expert <= 8u &&
            ds4_gpu_stream_expert_cache_configured_budget() >= n_expert &&
            ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                         down_expert_bytes) &&
            ds4_gpu_stream_expert_cache_effective_cap(layer_index,
                                                      n_total_expert,
                                                      n_expert) != 0;
        int32_t stream_selected_ids[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        ds4_gpu_stream_expert_cache_entry *stream_entries[8] = {
            NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL
        };
        uint64_t stream_gate_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        uint64_t stream_up_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        uint64_t stream_down_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        uint32_t stream_missing_mask = 0;
        uint32_t stream_entry_count = 0;
        uint32_t stream_resident_mask = 0;
        BOOL use_stream_split_deferred = false;
        id<MTLBuffer> stream_gate_addr_buf = nil;
        id<MTLBuffer> stream_up_addr_buf = nil;
        id<MTLBuffer> stream_down_addr_buf = nil;
        const int glm_stream_timing =
            ds4_gpu_stream_expert_timing_summary_enabled();

        if (use_stream_expert_addr_table) {
            const int had_batch = g_batch_cb != nil;
            int stream_ok = 1;
            const int have_prefetched_selected =
                ds4_gpu_glm_stream_selected_prefetch_take(model_map,
                                                          model_size,
                                                          layer_index,
                                                          n_total_expert,
                                                          n_expert,
                                                          gate_offset,
                                                          up_offset,
                                                          down_offset,
                                                          gate_expert_bytes,
                                                          down_expert_bytes,
                                                          stream_selected_ids);
            if (have_prefetched_selected) {
                if (had_batch && g_batch_has_work &&
                    g_stream_expert_pending_load.active &&
                    ds4_gpu_flush_commands() == 0) {
                    stream_ok = 0;
                }
            } else {
                if (had_batch && ds4_gpu_end_commands() == 0) {
                    stream_ok = 0;
                }
                if (stream_ok &&
                    ds4_gpu_tensor_read(selected,
                                        0,
                                        stream_selected_ids,
                                        (uint64_t)n_expert * sizeof(stream_selected_ids[0])) == 0) {
                    stream_ok = 0;
                }
            }
            for (uint32_t i = 0; stream_ok && i < n_expert; i++) {
                if (stream_selected_ids[i] < 0 ||
                    (uint32_t)stream_selected_ids[i] >= n_total_expert) {
                    fprintf(stderr,
                            "ds4: Metal GLM routed MoE selected expert id %d is outside 0..%u\n",
                            stream_selected_ids[i],
                            n_total_expert);
                    stream_ok = 0;
                }
            }
            if (stream_ok) {
                ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index,
                                                                  stream_selected_ids,
                                                                  n_expert);
                if (!ds4_gpu_moe_selected_hotlist_record(layer_index,
                                                         stream_selected_ids,
                                                         n_expert,
                                                         n_total_expert)) {
                    stream_ok = 0;
                }
            }
            if (stream_ok) {
                g_glm_stream_expert_addr_table_building++;
                for (uint32_t i = 0; stream_ok && i < n_expert; i++) {
                    const uint64_t expert_id = (uint64_t)(uint32_t)stream_selected_ids[i];
                    if (expert_id > UINT64_MAX / gate_expert_bytes ||
                        expert_id > UINT64_MAX / down_expert_bytes) {
                        fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n");
                        stream_ok = 0;
                        break;
                    }
                    const uint64_t gate_rel = expert_id * gate_expert_bytes;
                    const uint64_t down_rel = expert_id * down_expert_bytes;
                    if (gate_rel > UINT64_MAX - gate_offset ||
                        gate_rel > UINT64_MAX - up_offset ||
                        down_rel > UINT64_MAX - down_offset) {
                        fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n");
                        stream_ok = 0;
                        break;
                    }
                    stream_gate_abs_offsets[i] = gate_offset + gate_rel;
                    stream_up_abs_offsets[i] = up_offset + gate_rel;
                    stream_down_abs_offsets[i] = down_offset + down_rel;
                    stream_entries[i] =
                        ds4_gpu_stream_expert_cache_peek(model_map,
                                                         model_size,
                                                         layer_index,
                                                         (uint32_t)stream_selected_ids[i],
                                                         n_total_expert,
                                                         n_expert,
                                                         stream_gate_abs_offsets[i],
                                                         stream_up_abs_offsets[i],
                                                         stream_down_abs_offsets[i],
                                                         gate_expert_bytes,
                                                         down_expert_bytes);
                    if (!stream_entries[i]) {
                        stream_missing_mask |= 1u << i;
                    } else {
                        stream_resident_mask |= 1u << i;
                    }
                }
                if (stream_ok && glm_stream_timing) {
                    ds4_gpu_stream_expert_timing_note_cache_class(
                            stream_resident_mask,
                            stream_missing_mask);
                }
                use_stream_split_deferred =
                    stream_ok &&
                    getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_SPLIT") == NULL &&
                    stream_missing_mask != 0 &&
                    stream_resident_mask != 0 &&
                    ds4_gpu_stream_expert_split_worthwhile(stream_resident_mask,
                                                           stream_missing_mask) &&
                    ds4_gpu_stream_expert_split_ready() &&
                    ((gate_pair_q2 &&
                      g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline != nil) ||
                     (!gate_pair_q2 && !gate_pair_q5 &&
                      g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline != nil));
                if (use_stream_split_deferred) {
                    const ds4_gpu_stream_expert_table table = {
                        .model_map = model_map,
                        .model_size = model_size,
                        .layer = layer_index,
                        .n_total_expert = n_total_expert,
                        .gate_offset = gate_offset,
                        .up_offset = up_offset,
                        .down_offset = down_offset,
                        .gate_expert_bytes = gate_expert_bytes,
                        .down_expert_bytes = down_expert_bytes,
                    };
                    if (!ds4_gpu_stream_expert_cache_begin_selected_load(
                                &table,
                                stream_selected_ids,
                                n_expert)) {
                        stream_ok = 0;
                    }
                }
                if (stream_ok && stream_missing_mask != 0 &&
                    !use_stream_split_deferred &&
                    !ds4_gpu_stream_expert_cache_load_selected_missing(
                            model_map,
                            model_size,
                            layer_index,
                            stream_selected_ids,
                            n_total_expert,
                            n_expert,
                            stream_gate_abs_offsets,
                            stream_up_abs_offsets,
                            stream_down_abs_offsets,
                            gate_expert_bytes,
                            down_expert_bytes,
                            stream_missing_mask,
                            stream_entries)) {
                    fprintf(stderr,
                            "ds4: Metal GLM streaming expert cache failed to load "
                            "layer=%u missing=0x%x budget=%u\n",
                            layer_index,
                            stream_missing_mask,
                            ds4_gpu_stream_expert_cache_configured_budget());
                    stream_ok = 0;
                }
                for (uint32_t i = 0; stream_ok && i < n_expert; i++) {
                    ds4_gpu_stream_expert_cache_entry *entry = stream_entries[i];
                    if (!entry) {
                        if (use_stream_split_deferred &&
                            (stream_missing_mask & (1u << i)) != 0) {
                            continue;
                        }
                        stream_ok = 0;
                        break;
                    }
                    if (!ds4_gpu_stream_expert_cache_set_addr_slot(
                                layer_index,
                                (uint32_t)stream_selected_ids[i],
                                entry->gate_buffer,
                                entry->gate_inner,
                                entry->up_buffer,
                                entry->up_inner,
                                entry->down_buffer,
                                entry->down_inner)) {
                        stream_ok = 0;
                        break;
                    }
                    stream_entry_count++;
                }
                if (stream_ok &&
                    !ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
                                                              &stream_gate_addr_buf,
                                                              &stream_up_addr_buf,
                                                              &stream_down_addr_buf)) {
                    stream_ok = 0;
                }
                g_glm_stream_expert_addr_table_building--;
            }
            if (!have_prefetched_selected && had_batch &&
                ds4_gpu_begin_commands() == 0) {
                stream_ok = 0;
            }
            if (!stream_ok) return 0;
            ds4_gpu_stream_expert_cache_prune_layer(layer_index,
                                                    n_total_expert,
                                                    n_expert,
                                                    stream_selected_ids,
                                                    n_expert);
            ds4_gpu_stream_expert_cache_prune_global(layer_index,
                                                     stream_selected_ids,
                                                     n_expert);
        }

        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        uint64_t down_inner = 0;
        id<MTLBuffer> gatebuf = nil;
        id<MTLBuffer> upbuf = nil;
        id<MTLBuffer> downbuf = nil;
        if (!use_stream_expert_addr_table) {
            gatebuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                               gate_offset, gate_tensor_bytes,
                                               &gate_inner);
            upbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                             up_offset, up_tensor_bytes,
                                             &up_inner);
            downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                               down_offset, down_tensor_bytes,
                                               &down_inner);
            if (!gatebuf || !upbuf || !downbuf) return 0;
        }

        id<MTLComputePipelineState> pair_pipeline =
            (use_stream_split_deferred ?
             (gate_pair_q2 ?
              ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline,
                                   "kernel_glm_q2_K_addr_pair_swiglu2_f32_masked") :
              ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline,
                                   "kernel_glm_q4_K_addr_pair_swiglu_f32_masked")) :
             use_stream_expert_addr_table ?
             (gate_pair_q2 ?
              ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline,
                                   "kernel_glm_q2_K_addr_pair_swiglu2_f32") :
              ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline,
                                   "kernel_glm_q4_K_addr_pair_swiglu_f32")) :
             gate_pair_q2 ?
             ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline,
                                  "kernel_glm_q2_K_pair_swiglu_f32") :
             gate_pair_q5 ?
             ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline,
                                  "kernel_glm_q5_K_pair_swiglu_f32") :
             ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline,
                                  "kernel_glm_q4_K_pair_swiglu2_f32"));
        id<MTLComputePipelineState> down_pipeline =
            (use_stream_expert_addr_table ?
             (down_scalar_q2 ?
              ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline,
                                   "kernel_glm_q2_K_addr_down_f32") :
              ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline,
                                   "kernel_glm_q4_K_addr_down_f32")) :
             down_scalar_q2 ?
             ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline,
                                  "kernel_glm_q2_K_down_f32") :
             down_scalar_q4 ?
             ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline,
                                  "kernel_glm_q4_K_down_f32") :
             down_simd_q5 ?
             ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline,
                                  "kernel_glm_q5_K_down_f32") :
             ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline,
                                  "kernel_glm_q6_K_down_f32"));
        if (!pair_pipeline || !down_pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const bool glm_moe_stage_profile =
            g_batch_cb != nil &&
            ds4_gpu_stage_profile_enabled_for_layer(
                    "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE",
                    "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE_LAYER",
                    layer_index);
        const char *glm_moe_stage_filter =
            getenv("DS4_METAL_GLM_MOE_STAGE_PROFILE_FILTER");
        const char *glm_pair_path =
            use_stream_expert_addr_table ?
            (gate_pair_q2 ? "q2_stream_addr_swiglu" :
                            "q4_stream_addr_swiglu") :
            gate_pair_q2 ? "q2_scalar_swiglu" :
            gate_pair_q5 ? "q5_pair_simd_swiglu" : "q4_pair2_simd_swiglu";
        const char *glm_down_path =
            use_stream_expert_addr_table ?
            (down_scalar_q2 ? "q2_stream_addr_down" :
                              "q4_stream_addr_down_simd") :
            down_scalar_q2 ? "q2_down_simd" :
            down_scalar_q4 ? "q4_down_simd" :
            down_simd_q5 ? "q5_down_simd" : "q6_down_simd";
        double glm_moe_stage_t0 = 0.0;
        if (glm_moe_stage_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;
            glm_moe_stage_t0 = ds4_gpu_now_ms();
        }
        int ok = 1;
#define DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE(name) do { \
            if (ok && glm_moe_stage_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = 0; \
                } else { \
                    const char *stage_name = (name); \
                    const double now_ms = ds4_gpu_now_ms(); \
                    const int print_stage = \
                        !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
                        strstr(stage_name, glm_moe_stage_filter) != NULL; \
                    if (print_stage) { \
                        fprintf(stderr, \
                                "ds4: Metal GLM routed MoE one stage layer=%u tokens=1 experts=%u " \
                                "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \
                                layer_index, n_expert, \
                                ds4_gpu_metal_tensor_type_name(gate_type), \
                                ds4_gpu_metal_tensor_type_name(down_type), \
                                glm_pair_path, glm_down_path, \
                                stage_name, now_ms - glm_moe_stage_t0); \
                    } \
                    glm_moe_stage_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = 0; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb) ok = 0; \
                    } \
                } \
            } \
        } while (0)

        ds4_gpu_glm_routed_moe_args args = {
            .tp_rank = g_tp_split_rank,
            .tp_world = g_tp_split_world,
            .tp_expert_base = tp_expert_base_host,
            .in_dim = expert_in_dim,
            .mid_dim = expert_mid_dim,
            .out_dim = out_dim,
            .n_total_expert = n_total_expert,
            .n_expert_used = n_expert,
            .n_tokens = 1,
            .mid_token_stride = n_expert * expert_mid_dim,
            .down_type = down_type,
            .gate_expert_bytes = gate_expert_bytes,
            .gate_row_bytes = gate_row_bytes,
            .up_expert_bytes = up_expert_bytes,
            .up_row_bytes = up_row_bytes,
            .down_expert_bytes = down_expert_bytes,
            .down_row_bytes = down_row_bytes,
        };
        const NSUInteger pair_x_groups =
            gate_pair_q2 ? (use_stream_expert_addr_table ?
                            (NSUInteger)((expert_mid_dim + 1u) / 2u) :
                            (NSUInteger)((expert_mid_dim + 7u) / 8u)) :
            gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) :
            use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) :
            (NSUInteger)((expert_mid_dim + 1u) / 2u);
        const NSUInteger pair_threadgroup_bytes = 0u;
        const NSUInteger pair_threads = 64u;
        const NSUInteger down_x_groups =
            down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) :
            down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) :
            down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) :
            down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) :
            (NSUInteger)out_dim;
        const NSUInteger down_threadgroup_bytes =
            (down_scalar_q2 || down_simd) ? 0u : 256u * sizeof(float);
        const NSUInteger down_threads =
            (down_scalar_q2 || down_simd) ? 64u : 256u;
        if (use_stream_expert_addr_table &&
            !ds4_gpu_stream_expert_cache_mark_entries_inflight(
                    stream_entries,
                    use_stream_split_deferred ? n_expert : stream_entry_count,
                    use_stream_split_deferred ? stream_resident_mask : 0)) {
            return 0;
        }

        const int glm_stream_split_timing =
            use_stream_split_deferred && glm_stream_timing;
        double glm_stream_split_t0 =
            glm_stream_split_timing ? ds4_gpu_now_ms() : 0.0;
        double glm_stream_split_resident_ms = 0.0;
        double glm_stream_split_missing_load_ms = 0.0;
        double glm_stream_split_missing_wait_ms = 0.0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pair_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        if (use_stream_split_deferred) {
            [enc setBytes:&stream_resident_mask length:sizeof(stream_resident_mask) atIndex:1];
            [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2];
            [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
            [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
            [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6];
            [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7];
        } else {
            [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf
                    offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner
                   atIndex:1];
            [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf
                    offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner
                   atIndex:2];
            [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
            [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4];
            [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5];
            [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
        }
        if (use_stream_expert_addr_table) {
            const uint32_t use_count =
                use_stream_split_deferred ? n_expert : stream_entry_count;
            const uint32_t use_mask =
                use_stream_split_deferred ? stream_resident_mask : 0;
            for (uint32_t i = 0; i < use_count; i++) {
                if (use_mask != 0 && (use_mask & (1u << i)) == 0) continue;
                [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead];
                [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead];
            }
        }
        if (pair_threadgroup_bytes != 0u) {
            [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups,
                                              (NSUInteger)n_expert,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("pair");

        if (ok && use_stream_split_deferred) {
            id<MTLCommandBuffer> resident_cb = nil;
            const int resident_owned = owned;
            if (owned) {
                resident_cb = cb;
                [resident_cb commit];
                cb = nil;
            } else {
                ok = ds4_gpu_flush_commands();
            }
            if (glm_stream_split_timing) {
                const double now_ms = ds4_gpu_now_ms();
                glm_stream_split_resident_ms = now_ms - glm_stream_split_t0;
                glm_stream_split_t0 = now_ms;
            }
            if (ok) {
                ok = ds4_gpu_stream_expert_pending_load_finish(stream_entries);
            }
            if (glm_stream_split_timing) {
                const double now_ms = ds4_gpu_now_ms();
                glm_stream_split_missing_load_ms = now_ms - glm_stream_split_t0;
                glm_stream_split_t0 = now_ms;
            }
            if (ok) {
                ds4_gpu_stream_expert_cache_prune_layer(layer_index,
                                                        n_total_expert,
                                                        n_expert,
                                                        stream_selected_ids,
                                                        n_expert);
                ds4_gpu_stream_expert_cache_prune_global(layer_index,
                                                         stream_selected_ids,
                                                         n_expert);
                ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
                                                              &stream_gate_addr_buf,
                                                              &stream_up_addr_buf,
                                                              &stream_down_addr_buf);
            }
            if (ok) {
                if (resident_owned) {
                    ok = ds4_gpu_wait_command_buffer(
                            resident_cb,
                            "GLM streaming split resident pair");
                    ds4_gpu_stream_expert_cache_note_owned_completed();
                } else {
                    ok = ds4_gpu_wait_pending_command_buffers(
                            "GLM streaming split resident pair");
                }
            }
            if (glm_stream_split_timing) {
                const double now_ms = ds4_gpu_now_ms();
                glm_stream_split_missing_wait_ms = now_ms - glm_stream_split_t0;
                glm_stream_split_t0 = now_ms;
                ds4_gpu_stream_expert_timing_note_split(
                        stream_resident_mask,
                        stream_missing_mask,
                        glm_stream_split_resident_ms,
                        glm_stream_split_missing_load_ms +
                        glm_stream_split_missing_wait_ms);
                ds4_gpu_stream_expert_timing_note_split_missing_detail(
                        glm_stream_split_missing_load_ms,
                        0.0,
                        0.0,
                        0.0,
                        glm_stream_split_missing_wait_ms);
            }
            if (ok &&
                !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_entries,
                                                                   n_expert,
                                                                   stream_missing_mask)) {
                ok = 0;
            }
            if (ok) {
                cb = ds4_gpu_command_buffer(&owned);
                if (!cb) ok = 0;
            }
            if (ok) {
                enc = ds4_gpu_compute_encoder(cb);
                [enc setComputePipelineState:pair_pipeline];
                [enc setBytes:&args length:sizeof(args) atIndex:0];
                [enc setBytes:&stream_missing_mask length:sizeof(stream_missing_mask) atIndex:1];
                [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2];
                [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3];
                [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
                [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
                [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6];
                [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7];
                for (uint32_t i = 0; i < n_expert; i++) {
                    if ((stream_missing_mask & (1u << i)) == 0) continue;
                    [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead];
                    [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead];
                }
                if (pair_threadgroup_bytes != 0u) {
                    [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0];
                }
                [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups,
                                                      (NSUInteger)n_expert,
                                                      1)
                     threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)];
                ds4_gpu_end_compute_encoder(cb, enc);
            }
        }

        if (!ok) return 0;

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:down_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf
                offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner
               atIndex:1];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
        if (use_stream_expert_addr_table) {
            const uint32_t use_count =
                use_stream_split_deferred ? n_expert : stream_entry_count;
            for (uint32_t i = 0; i < use_count; i++) {
                [enc useResource:stream_entries[i]->down_buffer usage:MTLResourceUsageRead];
            }
        }
        if (down_threadgroup_bytes != 0u) {
            [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("down");

        if (!ok) return 0;
        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed MoE")) return 0;
#undef DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE
    }

    return 1;
}

static bool ds4_gpu_glm_grouped_moe_fast_default(void) {
    return !g_quality_mode;
}

static bool ds4_gpu_glm_routed_moe_batch_grouped_available(
        uint32_t gate_type,
        uint32_t up_type,
        uint32_t down_type,
        uint32_t n_expert,
        uint32_t n_tokens) {
    if (n_tokens < 32u ||
        !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) ||
        !ds4_gpu_glm_down_type_supported(down_type) ||
        ds4_gpu_mul_mm_id_map0_name(n_expert) == NULL) {
        return false;
    }
    const bool fast_default = ds4_gpu_glm_grouped_moe_fast_default();
    if (!fast_default) {
        return false;
    }
    if (n_tokens < 96u) return false;

    return ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)) != nil &&
           ds4_gpu_routed_mm_pipeline(gate_type) != nil &&
           ds4_gpu_routed_mm_pipeline(up_type) != nil &&
           ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) != nil;
}

static bool ds4_gpu_glm_grouped_moe_layer_enabled(uint32_t layer_index) {
    (void)layer_index;
    return true;
}

static int ds4_gpu_glm_routed_moe_batch_grouped_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                up_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                up_expert_bytes,
        uint64_t                up_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        uint32_t                layer_index,
        const ds4_gpu_tensor *x,
        uint32_t                n_tokens) {
    if (!ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
                                                         up_type,
                                                         down_type,
                                                         n_expert,
                                                         n_tokens)) {
        return 0;
    }
    if (n_expert > UINT32_MAX / n_tokens ||
        (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
        return 0;
    }

    const uint32_t pair_rows = n_tokens * n_expert;
    if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim ||
        (uint64_t)pair_rows > UINT64_MAX / out_dim ||
        (uint64_t)n_tokens > UINT64_MAX / expert_in_dim ||
        (uint64_t)n_tokens > UINT64_MAX / out_dim) {
        return 0;
    }

    const bool mid_f16 = true;
    const NSUInteger mm_id_threadgroup_bytes = 8192u;
    const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim;
    const uint64_t down_values = (uint64_t)pair_rows * out_dim;
    const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim;
    const uint64_t out_values = (uint64_t)n_tokens * out_dim;
    if (compact_mid_values > UINT64_MAX / sizeof(float) ||
        compact_mid_values > UINT64_MAX / (mid_f16 ? sizeof(uint16_t) : sizeof(float)) ||
        down_values > UINT64_MAX / sizeof(float) ||
        x_values > UINT64_MAX / sizeof(float) ||
        out_values > UINT64_MAX / sizeof(float)) {
        return 0;
    }

    const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float);
    const uint64_t mid_bytes = compact_mid_values * (mid_f16 ? sizeof(uint16_t) : sizeof(float));
    const uint64_t down_scratch_bytes = down_values * sizeof(float);
    const uint64_t x_bytes = x_values * sizeof(float);
    const uint64_t out_bytes = out_values * sizeof(float);
    const uint64_t selected_values = (uint64_t)n_tokens * n_expert;
    const uint64_t selected_bytes = selected_values * sizeof(int32_t);
    const uint64_t weights_bytes = selected_values * sizeof(float);
    if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes ||
        gate_scratch_bytes > NSUIntegerMax ||
        gate_scratch_bytes * 2ull > NSUIntegerMax ||
        down_scratch_bytes > NSUIntegerMax) {
        return 0;
    }

    uint32_t first_expert = 0;
    uint32_t n_bind_expert = 0;
    ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
    gate_offset += (uint64_t)first_expert * gate_expert_bytes;
    up_offset += (uint64_t)first_expert * up_expert_bytes;
    down_offset += (uint64_t)first_expert * down_expert_bytes;
    const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
    const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(mid) < mid_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes) {
            fprintf(stderr, "ds4: Metal GLM grouped routed MoE received undersized activation buffers\n");
            return 0;
        }
        if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer,
                                           &g_moe_gate_scratch_bytes,
                                           (NSUInteger)(gate_scratch_bytes * 2ull),
                                           "ds4_glm_moe_gate_up_scratch")) {
            return 0;
        }
        if (n_expert > 1 &&
            !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
                                           &g_moe_down_scratch_bytes,
                                           (NSUInteger)down_scratch_bytes,
                                           "ds4_glm_moe_down_scratch")) {
            return 0;
        }

        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        uint64_t down_inner = 0;
        id<MTLBuffer> gatebuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                         gate_offset, gate_tensor_bytes,
                                                         &gate_inner);
        id<MTLBuffer> upbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                       up_offset, up_tensor_bytes,
                                                       &up_inner);
        id<MTLBuffer> downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                         down_offset, down_tensor_bytes,
                                                         &down_inner);
        if (!gatebuf || !upbuf || !downbuf) return 0;

        id<MTLComputePipelineState> map_pipeline =
            ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert));
        id<MTLComputePipelineState> gate_pipeline = ds4_gpu_routed_mm_pipeline(gate_type);
        id<MTLComputePipelineState> up_pipeline = ds4_gpu_routed_mm_pipeline(up_type);
        id<MTLComputePipelineState> down_pipeline =
            ds4_gpu_routed_mm_f16_rhs_pipeline(down_type);
        if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) {
            return 0;
        }

        ds4_gpu_mul_mm_id_map_args map_args =
            ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens);
        ds4_gpu_mul_mm_id_args gate_args =
            ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                          gate_row_bytes, gate_expert_bytes,
                                          1, n_expert, n_tokens);
        ds4_gpu_mul_mm_id_args up_args =
            ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                          up_row_bytes, up_expert_bytes,
                                          1, n_expert, n_tokens);
        ds4_gpu_mul_mm_id_args down_args =
            ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert,
                                                    down_row_bytes, down_expert_bytes,
                                                    n_expert, n_expert, n_tokens,
                                                    mid_f16 ? sizeof(uint16_t) : sizeof(float));
        gate_args.tp_rank = g_tp_split_rank;
        gate_args.tp_world = g_tp_split_world;
        gate_args.tp_expert_base = (int32_t)first_expert;
        up_args.tp_rank = g_tp_split_rank;
        up_args.tp_world = g_tp_split_world;
        up_args.tp_expert_base = (int32_t)first_expert;
        down_args.tp_rank = g_tp_split_rank;
        down_args.tp_world = g_tp_split_world;
        down_args.tp_expert_base = (int32_t)first_expert;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const bool glm_moe_stage_profile = false;
        const char *glm_moe_stage_filter = NULL;
        double glm_moe_stage_t0 = 0.0;
        if (glm_moe_stage_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;
            glm_moe_stage_t0 = ds4_gpu_now_ms();
        }

        int ok = 1;
#define DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE(name) do { \
            if (ok && glm_moe_stage_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = 0; \
                } else { \
                    const char *stage_name = (name); \
                    const double now_ms = ds4_gpu_now_ms(); \
                    const int print_stage = \
                        !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
                        strstr(stage_name, glm_moe_stage_filter) != NULL; \
                    if (print_stage) { \
                        fprintf(stderr, \
                                "ds4: Metal GLM grouped routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \
                                "gate=%s down=%s mid=%s %s=%.3f ms\n", \
                                layer_index, n_tokens, pair_rows, n_expert, \
                                ds4_gpu_metal_tensor_type_name(gate_type), \
                                ds4_gpu_metal_tensor_type_name(down_type), \
                                mid_f16 ? "f16" : "f32", \
                                stage_name, now_ms - glm_moe_stage_t0); \
                    } \
                    glm_moe_stage_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = 0; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb) ok = 0; \
                    } \
                } \
            } \
        } while (0)

        ok = ds4_gpu_encode_mul_mm_id_map(cb,
                                           map_pipeline,
                                           &map_args,
                                           &gate_args,
                                           selectedbuf,
                                           ds4_gpu_tensor_offset(selected));
        DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("map");
        if (ok) {
            ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                       gate_pipeline,
                                                       &gate_args,
                                                       gatebuf,
                                                       (NSUInteger)gate_inner,
                                                       xbuf,
                                                       ds4_gpu_tensor_offset(x),
                                                       g_moe_gate_scratch_buffer,
                                                       0,
                                                       mm_id_threadgroup_bytes);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("gate");
        if (ok) {
            ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                       up_pipeline,
                                                       &up_args,
                                                       upbuf,
                                                       (NSUInteger)up_inner,
                                                       xbuf,
                                                       ds4_gpu_tensor_offset(x),
                                                       g_moe_gate_scratch_buffer,
                                                       (NSUInteger)gate_scratch_bytes,
                                                       mm_id_threadgroup_bytes);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("up");
        if (ok) {
            ok = ds4_gpu_encode_moe_swiglu_weight(cb,
                                                   g_moe_gate_scratch_buffer,
                                                   0,
                                                   g_moe_gate_scratch_buffer,
                                                   (NSUInteger)gate_scratch_bytes,
                                                   midbuf,
                                                   ds4_gpu_tensor_offset(mid),
                                                   weightsbuf,
                                                   ds4_gpu_tensor_offset(weights),
                                                   expert_mid_dim,
                                                   pair_rows,
                                                   0.0f,
                                                   mid_f16);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("activation_weight");

        id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer;
        NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0;
        if (ok) {
            ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                       down_pipeline,
                                                       &down_args,
                                                       downbuf,
                                                       (NSUInteger)down_inner,
                                                       midbuf,
                                                       ds4_gpu_tensor_offset(mid),
                                                       down_dst,
                                                       down_dst_off,
                                                       mm_id_threadgroup_bytes);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("down");
        if (ok && n_expert > 1) {
            ok = ds4_gpu_encode_moe_sum_experts(cb,
                                                 down_dst,
                                                 down_dst_off,
                                                 outbuf,
                                                 ds4_gpu_tensor_offset(out),
                                                 out_dim,
                                                 n_expert,
                                                 n_tokens);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("sum");
        if (!ok) return 0;

        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM grouped routed batch MoE")) {
            return 0;
        }
#undef DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE
    }

    return 1;
}

static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *mid,
        id<MTLBuffer>          gate_addrs,
        id<MTLBuffer>          up_addrs,
        id<MTLBuffer>          down_addrs,
        id<MTLBuffer>          overflow_gate,
        id<MTLBuffer>          overflow_up,
        id<MTLBuffer>          overflow_down,
        uint32_t               gate_type,
        uint32_t               up_type,
        uint32_t               down_type,
        uint64_t               gate_expert_bytes,
        uint64_t               gate_row_bytes,
        uint64_t               up_expert_bytes,
        uint64_t               up_row_bytes,
        uint64_t               down_expert_bytes,
        uint64_t               down_row_bytes,
        uint32_t               expert_in_dim,
        uint32_t               expert_mid_dim,
        uint32_t               out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t               n_total_expert,
        uint32_t               n_expert,
        uint32_t               layer_index,
        const ds4_gpu_tensor *x,
        uint32_t               n_tokens,
        ds4_gpu_stream_expert_cache_entry * const *resources,
        uint32_t               resource_count) {
    if (n_expert > UINT32_MAX / n_tokens ||
        !gate_addrs || !up_addrs || !down_addrs ||
        !ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
                                                        up_type,
                                                        down_type,
                                                        n_expert,
                                                        n_tokens)) {
        return 0;
    }

    const uint32_t pair_rows = n_tokens * n_expert;
    if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim ||
        (uint64_t)pair_rows > UINT64_MAX / out_dim ||
        (uint64_t)n_tokens > UINT64_MAX / expert_in_dim ||
        (uint64_t)n_tokens > UINT64_MAX / out_dim) {
        return 0;
    }

    const bool mid_f16 = true;
    const NSUInteger mm_id_threadgroup_bytes = 8192u;
    const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim;
    const uint64_t down_values = (uint64_t)pair_rows * out_dim;
    const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim;
    const uint64_t out_values = (uint64_t)n_tokens * out_dim;
    if (compact_mid_values > UINT64_MAX / sizeof(float) ||
        compact_mid_values > UINT64_MAX / sizeof(uint16_t) ||
        down_values > UINT64_MAX / sizeof(float) ||
        x_values > UINT64_MAX / sizeof(float) ||
        out_values > UINT64_MAX / sizeof(float)) {
        return 0;
    }

    const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float);
    const uint64_t mid_bytes = compact_mid_values * sizeof(uint16_t);
    const uint64_t down_scratch_bytes = down_values * sizeof(float);
    const uint64_t x_bytes = x_values * sizeof(float);
    const uint64_t out_bytes = out_values * sizeof(float);
    const uint64_t selected_values = (uint64_t)n_tokens * n_expert;
    const uint64_t selected_bytes = selected_values * sizeof(int32_t);
    const uint64_t weights_bytes = selected_values * sizeof(float);
    if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes ||
        gate_scratch_bytes > NSUIntegerMax ||
        gate_scratch_bytes * 2ull > NSUIntegerMax ||
        down_scratch_bytes > NSUIntegerMax) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(mid) < mid_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes) {
            fprintf(stderr, "ds4: Metal GLM grouped-address routed MoE received undersized activation buffers\n");
            return 0;
        }
        if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer,
                                           &g_moe_gate_scratch_bytes,
                                           (NSUInteger)(gate_scratch_bytes * 2ull),
                                           "ds4_glm_moe_gate_up_scratch")) {
            return 0;
        }
        if (n_expert > 1 &&
            !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
                                           &g_moe_down_scratch_bytes,
                                           (NSUInteger)down_scratch_bytes,
                                           "ds4_glm_moe_down_scratch")) {
            return 0;
        }

        id<MTLComputePipelineState> map_pipeline =
            ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert));
        id<MTLComputePipelineState> gate_pipeline =
            ds4_gpu_routed_mm_addr_pipeline(gate_type);
        id<MTLComputePipelineState> up_pipeline =
            ds4_gpu_routed_mm_addr_pipeline(up_type);
        id<MTLComputePipelineState> down_pipeline =
            ds4_gpu_routed_mm_addr_f16_rhs_pipeline(down_type);
        if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) {
            return 0;
        }

        ds4_gpu_mul_mm_id_map_args map_args =
            ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens);
        ds4_gpu_mul_mm_id_args gate_args =
            ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                          gate_row_bytes, gate_expert_bytes,
                                          1, n_expert, n_tokens);
        ds4_gpu_mul_mm_id_args up_args =
            ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                          up_row_bytes, up_expert_bytes,
                                          1, n_expert, n_tokens);
        ds4_gpu_mul_mm_id_args down_args =
            ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert,
                                                    down_row_bytes, down_expert_bytes,
                                                    n_expert, n_expert, n_tokens,
                                                    sizeof(uint16_t));
        if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(resources,
                                                               resource_count,
                                                               0)) {
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const bool glm_moe_stage_profile = false;
        const char *glm_moe_stage_filter = NULL;
        double glm_moe_stage_t0 = 0.0;
        if (glm_moe_stage_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;
            glm_moe_stage_t0 = ds4_gpu_now_ms();
        }

        int ok = 1;
#define DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE(name) do { \
            if (ok && glm_moe_stage_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = 0; \
                } else { \
                    const char *stage_name = (name); \
                    const double now_ms = ds4_gpu_now_ms(); \
                    const int print_stage = \
                        !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
                        strstr(stage_name, glm_moe_stage_filter) != NULL; \
                    if (print_stage) { \
                        fprintf(stderr, \
                                "ds4: Metal GLM grouped-address routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \
                                "gate=%s down=%s mid=f16 %s=%.3f ms\n", \
                                layer_index, n_tokens, pair_rows, n_expert, \
                                ds4_gpu_metal_tensor_type_name(gate_type), \
                                ds4_gpu_metal_tensor_type_name(down_type), \
                                stage_name, now_ms - glm_moe_stage_t0); \
                    } \
                    glm_moe_stage_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = 0; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb) ok = 0; \
                    } \
                } \
            } \
        } while (0)

        ok = ds4_gpu_encode_mul_mm_id_map(cb,
                                           map_pipeline,
                                           &map_args,
                                           &gate_args,
                                           selectedbuf,
                                           ds4_gpu_tensor_offset(selected));
        DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("map");
        if (ok) {
            ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb,
                                                           gate_pipeline,
                                                           &gate_args,
                                                           gate_addrs,
                                                           xbuf,
                                                           ds4_gpu_tensor_offset(x),
                                                           g_moe_gate_scratch_buffer,
                                                           0,
                                                           mm_id_threadgroup_bytes,
                                                           resources,
                                                           resource_count,
                                                           0,
                                                           overflow_gate);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("gate");
        if (ok) {
            ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb,
                                                           up_pipeline,
                                                           &up_args,
                                                           up_addrs,
                                                           xbuf,
                                                           ds4_gpu_tensor_offset(x),
                                                           g_moe_gate_scratch_buffer,
                                                           (NSUInteger)gate_scratch_bytes,
                                                           mm_id_threadgroup_bytes,
                                                           resources,
                                                           resource_count,
                                                           1,
                                                           overflow_up);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("up");
        if (ok) {
            ok = ds4_gpu_encode_moe_swiglu_weight(cb,
                                                   g_moe_gate_scratch_buffer,
                                                   0,
                                                   g_moe_gate_scratch_buffer,
                                                   (NSUInteger)gate_scratch_bytes,
                                                   midbuf,
                                                   ds4_gpu_tensor_offset(mid),
                                                   weightsbuf,
                                                   ds4_gpu_tensor_offset(weights),
                                                   expert_mid_dim,
                                                   pair_rows,
                                                   0.0f,
                                                   mid_f16);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("activation_weight");

        id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer;
        NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0;
        if (ok) {
            ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb,
                                                           down_pipeline,
                                                           &down_args,
                                                           down_addrs,
                                                           midbuf,
                                                           ds4_gpu_tensor_offset(mid),
                                                           down_dst,
                                                           down_dst_off,
                                                           mm_id_threadgroup_bytes,
                                                           resources,
                                                           resource_count,
                                                           2,
                                                           overflow_down);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("down");
        if (ok && n_expert > 1) {
            ok = ds4_gpu_encode_moe_sum_experts(cb,
                                                 down_dst,
                                                 down_dst_off,
                                                 outbuf,
                                                 ds4_gpu_tensor_offset(out),
                                                 out_dim,
                                                 n_expert,
                                                 n_tokens);
        }
        DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("sum");
        if (!ok) return 0;

        if (!ds4_gpu_finish_command_buffer(cb, owned,
                                           "GLM grouped-address routed batch MoE")) {
            return 0;
        }
#undef DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE
    }

    return 1;
}

static int ds4_gpu_glm_routed_moe_batch_tensor_impl(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                up_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                up_expert_bytes,
        uint64_t                up_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        uint32_t                layer_index,
        const ds4_gpu_tensor *x,
        uint32_t                n_tokens,
        uint32_t                mid_token_stride,
        bool                    allow_grouped,
        bool                    force_scalar_q4_pair) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !mid || !model_map || !selected || !weights || !x ||
        n_tokens == 0 ||
        n_total_expert == 0 || n_expert == 0 || n_expert > 256u ||
        n_expert > n_total_expert ||
        expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 ||
        gate_expert_bytes == 0 || gate_row_bytes == 0 ||
        up_expert_bytes == 0 || up_row_bytes == 0 ||
        down_expert_bytes == 0 || down_row_bytes == 0 ||
        (expert_in_dim % 256u) != 0 ||
        (expert_mid_dim % 256u) != 0 ||
        !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) ||
        !ds4_gpu_glm_down_type_supported(down_type)) {
        return 0;
    }

    const uint64_t per_token_mid = (uint64_t)n_expert * expert_mid_dim;
    if ((uint64_t)mid_token_stride < per_token_mid ||
        (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes ||
        (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes ||
        (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes ||
        (uint64_t)out_dim > UINT64_MAX / down_row_bytes ||
        (uint64_t)n_tokens > UINT64_MAX / expert_in_dim ||
        (uint64_t)n_tokens > UINT64_MAX / out_dim) {
        fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor byte size overflow\n");
        return 0;
    }

    const uint64_t full_gate_tensor_bytes =
        (uint64_t)n_total_expert * gate_expert_bytes;
    const uint64_t full_up_tensor_bytes =
        (uint64_t)n_total_expert * up_expert_bytes;
    const uint64_t full_down_tensor_bytes =
        (uint64_t)n_total_expert * down_expert_bytes;
    if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes ||
        up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes ||
        down_expert_bytes != (uint64_t)out_dim * down_row_bytes) {
        fprintf(stderr, "ds4: Metal GLM routed batch MoE received inconsistent expert strides\n");
        return 0;
    }
    if (gate_offset > model_size || full_gate_tensor_bytes > model_size - gate_offset ||
        up_offset > model_size || full_up_tensor_bytes > model_size - up_offset ||
        down_offset > model_size || full_down_tensor_bytes > model_size - down_offset) {
        fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor range is outside the mapped model\n");
        return 0;
    }

    if (allow_grouped &&
        (!g_ssd_streaming_mode ||
         ds4_gpu_glm_streaming_prefill_full_layer_active()) &&
        ds4_gpu_glm_grouped_moe_layer_enabled(layer_index) &&
        ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
                                                       up_type,
                                                       down_type,
                                                       n_expert,
                                                       n_tokens)) {
        return ds4_gpu_glm_routed_moe_batch_grouped_tensor(out,
                                                           mid,
                                                           model_map,
                                                           model_size,
                                                           gate_offset,
                                                           up_offset,
                                                           down_offset,
                                                           gate_type,
                                                           up_type,
                                                           down_type,
                                                           gate_expert_bytes,
                                                           gate_row_bytes,
                                                           up_expert_bytes,
                                                           up_row_bytes,
                                                           down_expert_bytes,
                                                           down_row_bytes,
                                                           expert_in_dim,
                                                           expert_mid_dim,
                                                           out_dim,
                                                           selected,
                                                           weights,
                                                           n_total_expert,
                                                           n_expert,
                                                           layer_index,
                                                           x,
                                                           n_tokens);
    }

    uint32_t first_expert = 0;
    uint32_t n_bind_expert = 0;
    ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
    gate_offset += (uint64_t)first_expert * gate_expert_bytes;
    up_offset += (uint64_t)first_expert * up_expert_bytes;
    down_offset += (uint64_t)first_expert * down_expert_bytes;
    const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
    const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim;
        const uint64_t out_values = (uint64_t)n_tokens * out_dim;
        const uint64_t mid_values =
            (uint64_t)(n_tokens - 1u) * mid_token_stride + per_token_mid;
        const uint64_t selected_values = (uint64_t)n_tokens * n_expert;
        if (x_values > UINT64_MAX / sizeof(float) ||
            out_values > UINT64_MAX / sizeof(float) ||
            mid_values > UINT64_MAX / sizeof(float) ||
            selected_values > UINT64_MAX / sizeof(int32_t)) {
            fprintf(stderr, "ds4: Metal GLM routed batch MoE activation byte size overflow\n");
            return 0;
        }
        const uint64_t x_bytes = x_values * sizeof(float);
        const uint64_t mid_bytes = mid_values * sizeof(float);
        const uint64_t out_bytes = out_values * sizeof(float);
        const uint64_t selected_bytes = selected_values * sizeof(int32_t);
        const uint64_t weights_bytes = selected_values * sizeof(float);
        if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(mid) < mid_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes) {
            fprintf(stderr, "ds4: Metal GLM routed batch MoE received undersized activation buffers\n");
            return 0;
        }

        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        uint64_t down_inner = 0;
        id<MTLBuffer> gatebuf = nil;
        id<MTLBuffer> upbuf = nil;
        id<MTLBuffer> downbuf = nil;
        id<MTLBuffer> stream_gate_addr_buf = nil;
        id<MTLBuffer> stream_up_addr_buf = nil;
        id<MTLBuffer> stream_down_addr_buf = nil;
        id<MTLBuffer> stream_overflow_gate = nil;
        id<MTLBuffer> stream_overflow_up = nil;
        id<MTLBuffer> stream_overflow_down = nil;
        ds4_gpu_stream_expert_cache_entry
            *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL };
        uint32_t stream_resource_count = 0;
        uint32_t stream_unique = 0;

        const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K;
        const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K;
        const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K;
        const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K;
        const BOOL down_simd_q4 = down_scalar_q4;
        const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K;
        const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K;
        const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6;
        const BOOL stream_addr_q2 =
            gate_pair_q2 && down_scalar_q2 &&
            g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil &&
            g_glm_q2_k_addr_down_f32_pipeline != nil;
        const BOOL stream_addr_q4 =
            !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 &&
            g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil &&
            g_glm_q4_k_addr_down_f32_pipeline != nil;
        BOOL use_stream_expert_addr_table =
            g_ssd_streaming_mode &&
            !ds4_gpu_glm_streaming_prefill_full_layer_active() &&
            n_tokens > 1 &&
            (stream_addr_q2 || stream_addr_q4) &&
            layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
            n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
            n_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED &&
            ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                         down_expert_bytes) &&
            ds4_gpu_stream_expert_cache_effective_cap(layer_index,
                                                      n_total_expert,
                                                      n_expert) != 0;
        const BOOL use_stream_grouped_addr_table =
            use_stream_expert_addr_table &&
            allow_grouped &&
            getenv("DS4_METAL_GLM_DISABLE_STREAMING_GROUPED_ADDR_PREFILL") == NULL &&
            ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
                                                           up_type,
                                                           down_type,
                                                           n_expert,
                                                           n_tokens);
        const BOOL enable_q4_pair4 = true;
        const BOOL q4_scalar_pair = false;
        const BOOL q4_pair2 =
            !use_stream_expert_addr_table &&
            !gate_pair_q5 && !q4_scalar_pair &&
            (force_scalar_q4_pair || !enable_q4_pair4);
        id<MTLComputePipelineState> pair_pipeline =
            use_stream_expert_addr_table ?
             (gate_pair_q2 ?
              ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline,
                                   "kernel_glm_q2_K_addr_pair_swiglu2_f32") :
              ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline,
                                   "kernel_glm_q4_K_addr_pair_swiglu_f32")) :
            gate_pair_q2 ?
             ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline,
                                  "kernel_glm_q2_K_pair_swiglu_f32") :
            gate_pair_q5 ?
             ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline,
                                  "kernel_glm_q5_K_pair_swiglu_f32") :
            (q4_scalar_pair ?
             ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu_f32_pipeline,
                                  "kernel_glm_q4_K_pair_swiglu_f32") :
             q4_pair2 ?
              ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline,
                                   "kernel_glm_q4_K_pair_swiglu2_f32") :
              ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu4_f32_pipeline,
                                   "kernel_glm_q4_K_pair_swiglu4_f32"));
        id<MTLComputePipelineState> down_pipeline =
            use_stream_expert_addr_table ?
             (down_scalar_q2 ?
              ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline,
                                   "kernel_glm_q2_K_addr_down_f32") :
              ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline,
                                   "kernel_glm_q4_K_addr_down_f32")) :
            down_scalar_q2 ?
            ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline,
                                 "kernel_glm_q2_K_down_f32") :
            down_scalar_q4 ?
            ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline,
                                 "kernel_glm_q4_K_down_f32") :
            down_simd_q5 ?
            ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline,
                                 "kernel_glm_q5_K_down_f32") :
            ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline,
                                 "kernel_glm_q6_K_down_f32");
        if (!pair_pipeline || !down_pipeline) return 0;

        if (use_stream_expert_addr_table) {
            const int had_batch = g_batch_cb != nil;
            if (had_batch && ds4_gpu_end_commands() == 0) {
                return 0;
            }
            if (!ds4_gpu_stream_expert_cache_prepare_selected_batch(
                        model_map,
                        model_size,
                        layer_index,
                        selected,
                        n_tokens,
                        n_total_expert,
                        n_expert,
                        gate_offset,
                        up_offset,
                        down_offset,
                        gate_expert_bytes,
                        down_expert_bytes,
                        &stream_gate_addr_buf,
                        &stream_up_addr_buf,
                        &stream_down_addr_buf,
                        stream_resources,
                        &stream_resource_count,
                        &stream_unique,
                        &stream_overflow_gate,
                        &stream_overflow_up,
                        &stream_overflow_down)) {
                return 0;
            }
            if (stream_unique == 0) {
                ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                return 0;
            }
            if (had_batch && ds4_gpu_begin_commands() == 0) {
                ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                return 0;
            }
            if (use_stream_grouped_addr_table) {
                return ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor(
                        out,
                        mid,
                        stream_gate_addr_buf,
                        stream_up_addr_buf,
                        stream_down_addr_buf,
                        stream_overflow_gate,
                        stream_overflow_up,
                        stream_overflow_down,
                        gate_type,
                        up_type,
                        down_type,
                        gate_expert_bytes,
                        gate_row_bytes,
                        up_expert_bytes,
                        up_row_bytes,
                        down_expert_bytes,
                        down_row_bytes,
                        expert_in_dim,
                        expert_mid_dim,
                        out_dim,
                        selected,
                        weights,
                        n_total_expert,
                        n_expert,
                        layer_index,
                        x,
                        n_tokens,
                        stream_resources,
                        stream_resource_count);
            }
        } else {
            gatebuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                               gate_offset, gate_tensor_bytes,
                                               &gate_inner);
            upbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                             up_offset, up_tensor_bytes,
                                             &up_inner);
            downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                               down_offset, down_tensor_bytes,
                                               &down_inner);
            if (!gatebuf || !upbuf || !downbuf) return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        const bool glm_moe_stage_profile = false;
        const char *glm_moe_stage_filter = NULL;
        const char *glm_pair_path = use_stream_expert_addr_table ?
                                    (gate_pair_q2 ? "q2_stream_addr_swiglu" :
                                                    "q4_stream_addr_swiglu") :
                                    gate_pair_q2 ? "q2_scalar_swiglu" :
                                    gate_pair_q5 ? "q5_pair_simd_swiglu" :
                                    (q4_scalar_pair ? "q4_scalar_swiglu" :
                                    (q4_pair2 ? "q4_pair2_simd_swiglu" :
                                                "q4_pair4_simd_swiglu"));
        const char *glm_down_path =
            use_stream_expert_addr_table ?
            (down_scalar_q2 ? "q2_stream_addr_down" : "q4_stream_addr_down_simd") :
            down_scalar_q2 ? "q2_down_scalar" :
            down_scalar_q4 ? "q4_down_simd" :
            down_simd_q5 ? "q5_down_simd" : "q6_down_simd";
        double glm_moe_stage_t0 = 0.0;
        if (glm_moe_stage_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;
            glm_moe_stage_t0 = ds4_gpu_now_ms();
        }
        int ok = 1;
#define DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE(name) do { \
            if (ok && glm_moe_stage_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = 0; \
                } else { \
                    const char *stage_name = (name); \
                    const double now_ms = ds4_gpu_now_ms(); \
                    const int print_stage = \
                        !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
                        strstr(stage_name, glm_moe_stage_filter) != NULL; \
                    if (print_stage) { \
                        fprintf(stderr, \
                                "ds4: Metal GLM routed MoE batch stage layer=%u tokens=%u experts=%u " \
                                "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \
                                layer_index, n_tokens, n_expert, \
                                ds4_gpu_metal_tensor_type_name(gate_type), \
                                ds4_gpu_metal_tensor_type_name(down_type), \
                                glm_pair_path, glm_down_path, \
                                stage_name, now_ms - glm_moe_stage_t0); \
                    } \
                    glm_moe_stage_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = 0; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb) ok = 0; \
                    } \
                } \
            } \
        } while (0)

        ds4_gpu_glm_routed_moe_args args = {
            .tp_rank = g_tp_split_rank,
            .tp_world = g_tp_split_world,
            .tp_expert_base = (int32_t)first_expert,
            .in_dim = expert_in_dim,
            .mid_dim = expert_mid_dim,
            .out_dim = out_dim,
            .n_total_expert = n_total_expert,
            .n_expert_used = n_expert,
            .n_tokens = n_tokens,
            .mid_token_stride = mid_token_stride,
            .down_type = down_type,
            .gate_expert_bytes = gate_expert_bytes,
            .gate_row_bytes = gate_row_bytes,
            .up_expert_bytes = up_expert_bytes,
            .up_row_bytes = up_row_bytes,
            .down_expert_bytes = down_expert_bytes,
            .down_row_bytes = down_row_bytes,
        };
        const NSUInteger pair_x_groups =
            gate_pair_q2 ? (use_stream_expert_addr_table ?
                            (NSUInteger)((expert_mid_dim + 1u) / 2u) :
                            (NSUInteger)((expert_mid_dim + 7u) / 8u)) :
            gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) :
            use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) :
            q4_scalar_pair ? (NSUInteger)expert_mid_dim :
            q4_pair2 ? (NSUInteger)((expert_mid_dim + 1u) / 2u) :
            (NSUInteger)((expert_mid_dim + 7u) / 8u);
        const NSUInteger pair_threadgroup_bytes =
            q4_scalar_pair ? 512u * sizeof(float) : 0u;
        const NSUInteger pair_threads =
            q4_scalar_pair ? 256u : 64u;
        const NSUInteger down_x_groups =
            down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) :
            down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) :
            down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) :
            down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) :
            (NSUInteger)out_dim;
        const NSUInteger down_threadgroup_bytes =
            down_simd ? 0u : 256u * sizeof(float);
        const NSUInteger down_threads =
            down_simd ? 64u : 256u;
        if (use_stream_expert_addr_table &&
            !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_resources,
                                                               stream_resource_count,
                                                               0)) {
            return 0;
        }

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pair_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf
                offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner
               atIndex:1];
        [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf
                offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner
               atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4];
        [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
        if (use_stream_expert_addr_table) {
            for (uint32_t i = 0; i < stream_resource_count; i++) {
                [enc useResource:stream_resources[i]->gate_buffer usage:MTLResourceUsageRead];
                [enc useResource:stream_resources[i]->up_buffer usage:MTLResourceUsageRead];
            }
            if (stream_overflow_gate) [enc useResource:stream_overflow_gate usage:MTLResourceUsageRead];
            if (stream_overflow_up) [enc useResource:stream_overflow_up usage:MTLResourceUsageRead];
        }
        if (pair_threadgroup_bytes != 0u) {
            [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups,
                                              (NSUInteger)n_expert,
                                              (NSUInteger)n_tokens)
             threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("pair");

        enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:down_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf
                offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner
               atIndex:1];
        [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
        if (use_stream_expert_addr_table) {
            for (uint32_t i = 0; i < stream_resource_count; i++) {
                [enc useResource:stream_resources[i]->down_buffer usage:MTLResourceUsageRead];
            }
            if (stream_overflow_down) [enc useResource:stream_overflow_down usage:MTLResourceUsageRead];
        }
        if (down_threadgroup_bytes != 0u) {
            [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0];
        }
        [enc dispatchThreadgroups:MTLSizeMake(down_x_groups,
                                              (NSUInteger)n_tokens,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);
        DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("down");

        if (!ok) return 0;
        if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed batch MoE")) return 0;
#undef DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE
    }

    return 1;
}

int ds4_gpu_glm_routed_moe_batch_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                up_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                up_expert_bytes,
        uint64_t                up_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        uint32_t                layer_index,
        const ds4_gpu_tensor *x,
        uint32_t                n_tokens,
        uint32_t                mid_token_stride,
        bool                    force_resident) {
    (void)force_resident;
    return ds4_gpu_glm_routed_moe_batch_tensor_impl(out,
                                                    mid,
                                                    model_map,
                                                    model_size,
                                                    gate_offset,
                                                    up_offset,
                                                    down_offset,
                                                    gate_type,
                                                    up_type,
                                                    down_type,
                                                    gate_expert_bytes,
                                                    gate_row_bytes,
                                                    up_expert_bytes,
                                                    up_row_bytes,
                                                    down_expert_bytes,
                                                    down_row_bytes,
                                                    expert_in_dim,
                                                    expert_mid_dim,
                                                    out_dim,
                                                    selected,
                                                    weights,
                                                    n_total_expert,
                                                    n_expert,
                                                    layer_index,
                                                    x,
                                                    n_tokens,
                                                    mid_token_stride,
                                                    true,
                                                    false);
}

int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                up_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                up_expert_bytes,
        uint64_t                up_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        uint32_t                layer_index,
        const ds4_gpu_tensor *x,
        uint32_t                n_tokens,
        uint32_t                mid_token_stride) {
    return ds4_gpu_glm_routed_moe_batch_tensor_impl(out,
                                                    mid,
                                                    model_map,
                                                    model_size,
                                                    gate_offset,
                                                    up_offset,
                                                    down_offset,
                                                    gate_type,
                                                    up_type,
                                                    down_type,
                                                    gate_expert_bytes,
                                                    gate_row_bytes,
                                                    up_expert_bytes,
                                                    up_row_bytes,
                                                    down_expert_bytes,
                                                    down_row_bytes,
                                                    expert_in_dim,
                                                    expert_mid_dim,
                                                    out_dim,
                                                    selected,
                                                    weights,
                                                    n_total_expert,
                                                    n_expert,
                                                    layer_index,
                                                    x,
                                                    n_tokens,
                                                    mid_token_stride,
                                                    false,
                                                    false);
}

int ds4_gpu_router_select_tensor(
        ds4_gpu_tensor       *selected,
        ds4_gpu_tensor       *weights,
        ds4_gpu_tensor       *probs,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                bias_offset,
        uint64_t                hash_offset,
        uint32_t                hash_rows,
        uint32_t                token,
        uint32_t                n_expert,
        uint32_t                n_expert_used,
        float                   expert_weight_scale,
        uint32_t                n_expert_groups,
        uint32_t                n_group_used,
        bool                    has_bias,
        bool                    hash_mode,
        const ds4_gpu_tensor *logits) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || !weights || !probs || !logits || !model_map ||
        n_expert == 0 || n_expert_used == 0) return 0;
    if (hash_mode && token >= hash_rows) return 0;
    if (n_expert_groups > 1u || n_group_used > 0u) {
        fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n");
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
        if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf ||
            ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) ||
            ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int) ||
            ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) ||
            ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) {
            fprintf(stderr, "ds4: Metal router select received undersized buffers\n");
            return 0;
        }

        uint64_t bias_inner = 0;
        uint64_t hash_inner = 0;
        id<MTLBuffer> biasbuf = nil;
        id<MTLBuffer> hashbuf = nil;
        NSUInteger bias_set_offset = 0;
        NSUInteger hash_set_offset = 0;
        if (has_bias && !hash_mode) {
            const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
            biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner);
            if (!biasbuf) return 0;
            bias_set_offset = (NSUInteger)bias_inner;
        }
        if (hash_mode) {
            const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t);
            hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner);
            if (!hashbuf) return 0;
            hash_set_offset = (NSUInteger)hash_inner;
        }

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        const int32_t token_i32 = (int32_t)token;
        int ok = cb &&
                 ds4_gpu_encode_router_select(cb,
                                                      selected,
                                                      weights,
                                                      probs,
                                                      logitsbuf,
                                                      ds4_gpu_tensor_offset(logits),
                                                      biasbuf,
                                                      bias_set_offset,
                                                      hashbuf,
                                                      hash_set_offset,
                                                      nil,
                                                      0,
                                                      &token_i32,
                                                      hash_rows,
                                                      1,
                                                      n_expert,
                                                      n_expert_used,
                                                      expert_weight_scale,
                                                      has_bias && !hash_mode,
                                                      hash_mode);
        if (!had_batch) {
            ok = ds4_gpu_end_commands() != 0 && ok;
        }
        if (!ok) return 0;
    }

    return 1;
}

int ds4_gpu_router_select_batch_tensor(
        ds4_gpu_tensor       *selected,
        ds4_gpu_tensor       *weights,
        ds4_gpu_tensor       *probs,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                bias_offset,
        uint64_t                hash_offset,
        uint32_t                hash_rows,
        uint32_t                n_expert_groups,
        uint32_t                n_group_used,
        bool                    has_bias,
        bool                    hash_mode,
        const ds4_gpu_tensor *logits,
        const ds4_gpu_tensor *tokens,
        uint32_t                n_expert,
        uint32_t                n_expert_used,
        float                   expert_weight_scale,
        uint32_t                n_tokens) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!selected || !weights || !probs || !logits || !tokens || !model_map ||
        n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0;
    if (n_expert_groups > 1u || n_group_used > 0u) {
        fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n");
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
        id<MTLBuffer> tokensbuf = ds4_gpu_tensor_buffer(tokens);
        if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokensbuf ||
            ds4_gpu_tensor_bytes(logits) < (uint64_t)n_tokens * n_expert * sizeof(float) ||
            ds4_gpu_tensor_bytes(selected) < (uint64_t)n_tokens * n_expert_used * sizeof(int) ||
            ds4_gpu_tensor_bytes(weights) < (uint64_t)n_tokens * n_expert_used * sizeof(float) ||
            ds4_gpu_tensor_bytes(probs) < (uint64_t)n_tokens * n_expert * sizeof(float) ||
            ds4_gpu_tensor_bytes(tokens) < (uint64_t)n_tokens * sizeof(int32_t)) {
            fprintf(stderr, "ds4: Metal router batch select received undersized buffers\n");
            return 0;
        }

        uint64_t bias_inner = 0;
        uint64_t hash_inner = 0;
        id<MTLBuffer> biasbuf = nil;
        id<MTLBuffer> hashbuf = nil;
        NSUInteger bias_set_offset = 0;
        NSUInteger hash_set_offset = 0;
        if (has_bias && !hash_mode) {
            const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
            biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner);
            if (!biasbuf) return 0;
            bias_set_offset = (NSUInteger)bias_inner;
        }
        if (hash_mode) {
            const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t);
            hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner);
            if (!hashbuf) return 0;
            hash_set_offset = (NSUInteger)hash_inner;
        }

        const bool had_batch = g_batch_cb != nil;
        if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        int ok = cb &&
                 ds4_gpu_encode_router_select(cb,
                                                      selected,
                                                      weights,
                                                      probs,
                                                      logitsbuf,
                                                      ds4_gpu_tensor_offset(logits),
                                                      biasbuf,
                                                      bias_set_offset,
                                                      hashbuf,
                                                      hash_set_offset,
                                                      tokensbuf,
                                                      ds4_gpu_tensor_offset(tokens),
                                                      NULL,
                                                      hash_rows,
                                                      n_tokens,
                                                      n_expert,
                                                      n_expert_used,
                                                      expert_weight_scale,
                                                      has_bias && !hash_mode,
                                                      hash_mode);
        if (!had_batch) {
            ok = ds4_gpu_end_commands() != 0 && ok;
        }
        if (!ok) return 0;
    }

    return 1;
}

int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) {
    if (n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
        (!selected && n_selected != 0)) return 0;
    for (uint32_t i = 0; i < n_selected; i++) {
        g_routed_moe_selected_override[i] = selected[i];
    }
    g_routed_moe_selected_override_n = n_selected;
    return 1;
}

int ds4_gpu_routed_moe_one_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        ds4_gpu_tensor       *experts,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        float                   clamp,
        const ds4_gpu_tensor *x,
        const ds4_gpu_tensor *add_in,
        uint32_t                layer_index,
        bool                    force_resident) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    /* TP sharding: only the owned contiguous expert range is mapped,
     * so bind from the owned base, validate only its bytes, and tell the
     * kernels the first expert id present at that base. */
    uint32_t first_expert = 0;
    uint32_t n_bind_expert = 0;
    ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
    const int32_t tp_expert_base_host = (int32_t)first_expert;
    gate_offset += (uint64_t)first_expert * gate_expert_bytes;
    up_offset += (uint64_t)first_expert * gate_expert_bytes;
    down_offset += (uint64_t)first_expert * down_expert_bytes;

    if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights ||
        n_total_expert == 0 || n_expert == 0 ||
        n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
        gate_expert_bytes == 0 || down_expert_bytes == 0 ||
        gate_row_bytes == 0 || down_row_bytes == 0) {
        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32044);
        return 0;
    }
    if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32047); return 0; }
    ds4_gpu_stream_expert_cache_note_token(layer_index);

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
        id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> expertsbuf = ds4_gpu_tensor_buffer(experts);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> selected_exec_buf = selectedbuf;
        NSUInteger selected_exec_off = ds4_gpu_tensor_offset(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float);
        const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float);
        if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(gate) < mid_bytes ||
            ds4_gpu_tensor_bytes(up) < mid_bytes ||
            ds4_gpu_tensor_bytes(mid) < mid_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int) ||
            ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) {
            fprintf(stderr, "ds4: Metal routed tensor MoE received undersized activation buffers\n");
            return 0;
        }
        if (n_expert > 1 &&
            (!expertsbuf ||
             ds4_gpu_tensor_bytes(experts) < (uint64_t)n_expert * out_dim * sizeof(float))) {
            fprintf(stderr, "ds4: Metal routed tensor MoE received undersized expert output buffer\n");
            return 0;
        }

        if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
            (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
            fprintf(stderr, "ds4: Metal routed MoE tensor byte size overflow\n");
            return 0;
        }
        const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
        const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        uint64_t down_inner = 0;
        id<MTLBuffer> gate_buf = nil;
        id<MTLBuffer> up_buf = nil;
        id<MTLBuffer> down_buf = nil;
        __unsafe_unretained id<MTLBuffer> gate_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil };
        __unsafe_unretained id<MTLBuffer> up_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil };
        __unsafe_unretained id<MTLBuffer> down_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil };
        ds4_gpu_stream_expert_cache_entry *stream_slot_entries[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { NULL };
        NSUInteger gate_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
        NSUInteger up_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
        NSUInteger down_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
        uint64_t stream_gate_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
        uint64_t stream_up_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
        uint64_t stream_down_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
        id<MTLBuffer> stream_gate_addr_buf = nil;
        id<MTLBuffer> stream_up_addr_buf = nil;
        id<MTLBuffer> stream_down_addr_buf = nil;
        bool use_stream_expert_addr_table = false;
        bool use_stream_expert_masked_addr_table = false;
        bool use_stream_compact_addr_table = false;
        bool use_stream_expert_cache = false;
        bool use_stream_expert_split_candidate = false;
        bool use_stream_expert_split_deferred = false;
        bool stream_expert_split_completed = false;
        uint32_t stream_expert_resident_mask = 0;
        uint32_t stream_expert_missing_mask = 0;
        __unsafe_unretained id<MTLBuffer> gate_group6_bufs[6] = { nil, nil, nil, nil, nil, nil };
        __unsafe_unretained id<MTLBuffer> up_group6_bufs[6] = { nil, nil, nil, nil, nil, nil };
        __unsafe_unretained id<MTLBuffer> down_group6_bufs[6] = { nil, nil, nil, nil, nil, nil };
        NSUInteger gate_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        NSUInteger up_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        NSUInteger down_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        __unsafe_unretained id<MTLBuffer> gate_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil };
        __unsafe_unretained id<MTLBuffer> up_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil };
        __unsafe_unretained id<MTLBuffer> down_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil };
        NSUInteger gate_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
        NSUInteger up_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
        NSUInteger down_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
        __unsafe_unretained id<MTLBuffer> gate_group24_bufs[24] = { nil };
        __unsafe_unretained id<MTLBuffer> up_group24_bufs[24] = { nil };
        __unsafe_unretained id<MTLBuffer> down_group24_bufs[24] = { nil };
        NSUInteger gate_group24_offsets[24] = { 0 };
        NSUInteger up_group24_offsets[24] = { 0 };
        NSUInteger down_group24_offsets[24] = { 0 };
        DS4MetalQ4ExpertTable *gate_table = nil;
        DS4MetalQ4ExpertTable *up_table = nil;
        DS4MetalQ4ExpertTable *down_table = nil;
        id q4_table_layer_residency = nil;
        int32_t selected_ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };

        const uint32_t n_tokens = 1;
        const uint32_t pair_rows = n_tokens * n_expert;
        const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float);
        if ((n_expert > 1 && !expertsbuf &&
             !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
                                              &g_moe_down_scratch_bytes,
                                              (NSUInteger)down_scratch_bytes,
                                              "ds4_moe_down_scratch"))) {
            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32148);
            return 0;
        }

        const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type);
        const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type);
        id<MTLComputePipelineState> gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type);
        id<MTLComputePipelineState> down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type);
        if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) {
            fprintf(stderr, "ds4: unsupported Metal routed MoE quant types gate=%u down=%u\n",
                    gate_type, down_type);
            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32158);
            return 0;
        }

        ds4_gpu_mul_mv_id_args gate_args =
            ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                          gate_row_bytes, gate_expert_bytes,
                                          1, n_expert, n_tokens, gate_nr0);
        ds4_gpu_mul_mv_id_args down_args =
            ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert,
                                          down_row_bytes, down_expert_bytes,
                                          n_expert, n_expert, n_tokens, down_nr0);
        /* Tensor-parallel expert ownership; non-TP calls keep tp_world at 1. */
        gate_args.tp_rank = g_tp_split_rank;
        gate_args.tp_world = g_tp_split_world;
        gate_args.tp_expert_base = tp_expert_base_host;
        down_args.tp_rank = g_tp_split_rank;
        down_args.tp_world = g_tp_split_world;
        down_args.tp_addend = add_in != NULL;
        down_args.tp_expert_base = tp_expert_base_host;

        const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type);
        const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type);
        const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type);
        const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type);
        const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type);
        const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type);
        int ok = 1;
        const bool write_clamped_moe =
            getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL;
        id<MTLComputePipelineState> pair_swiglu_pipeline = nil;
        if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) {
            pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline;
        } else if (gate_type == DS4_METAL_TENSOR_Q4_K) {
            pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline;
        }
        const bool fuse_pair_swiglu =
            !g_quality_mode &&
            !write_clamped_moe &&
            getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
            pair_swiglu_pipeline != nil;
        id<MTLComputePipelineState> down_sum6_pipeline = nil;
        if (down_type == DS4_METAL_TENSOR_Q2_K) {
            down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline;
        } else if (down_type == DS4_METAL_TENSOR_Q4_K) {
            down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline;
        } else if (down_type == DS4_METAL_TENSOR_IQ2_XXS &&
                   g_tp_split_world == 2) {
            /* IQ2 down-sum exists for the GLM TP resident split only; the
             * streaming paths must keep their addr/masked chain (expert
             * bytes are not at their model-map offsets when streamed). */
            down_sum6_pipeline = g_moe_mul_mv_id_iq2_xxs_sum6_pipeline;
        }
        const bool direct_down_sum =
            !g_quality_mode &&
            (n_expert == 6 || (n_expert == 8 && g_tp_split_world == 2)) &&
            n_tokens == 1 &&
            down_sum6_pipeline != nil;
        /* The expert-ownership split lives only in the fused id pair+sum6
         * kernels; every other routed variant would silently compute full
         * sums on both ranks and double the combine.  Fail fast instead. */
        if ((g_tp_split_world > 1 || add_in) && !(fuse_pair_swiglu && direct_down_sum)) {
            fprintf(stderr, "ds4: tensor-parallel routed MoE requires the fused pair+sum6 decode path\n");
            return 0;
        }
        const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull;
        /*
         * The grouped Q4 experiment keeps selected IDs on GPU, but it also walks
         * every expert window in the layer.  On PRO Q4 this measured far slower
         * than the active selected-slot path, so keep it opt-in for profiling.
         */
        const bool enable_q4_grouped_experts =
            getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL &&
            getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL;
        const bool use_q4_grouped_experts =
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert >= 128 &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_group_q4_k_sum6_pipeline != nil &&
            enable_q4_grouped_experts;
        const uint32_t q4_expert_group_size =
            use_q4_grouped_experts ? ds4_gpu_q4_expert_group_size(n_total_expert) : 0;
        const bool q4_grouped_boundary =
            use_q4_grouped_experts &&
            g_batch_cb != nil &&
            getenv("DS4_METAL_DISABLE_Q4_GROUPED_BOUNDARY") == NULL;
        const bool q4_grouped_cache_views =
            getenv("DS4_METAL_Q4_GROUPED_CACHE_VIEWS") != NULL;
        const uint32_t q4_group6_expert_group_size = 64;
        const bool use_q4_group6_experts =
            !use_q4_grouped_experts &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == q4_group6_expert_group_size * 6u &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil &&
            getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL &&
            getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL;
        const uint32_t q4_group8_expert_group_size = 48;
        const bool use_q4_group8_experts =
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == q4_group8_expert_group_size * 8u &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil &&
            getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL &&
            getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL;
        const uint32_t q4_group24_expert_group_size = 16;
        const bool use_q4_group24_experts =
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            !use_q4_group8_experts &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == q4_group24_expert_group_size * 24u &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            direct_down_sum &&
            g_moe_mul_mv_group24_q4_k_id_pipeline != nil &&
            g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil &&
            getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL &&
            getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL;
        const bool q4_group24_exact_views =
            use_q4_group24_experts &&
            getenv("DS4_METAL_Q4_GROUP24_EXACT_VIEWS") != NULL &&
            getenv("DS4_METAL_Q4_GROUP24_BASE_VIEWS") == NULL;
        const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0;
        const bool can_wrap_q4_exact_tensors =
            max_buffer_len != 0 &&
            gate_tensor_bytes <= max_buffer_len &&
            down_tensor_bytes <= max_buffer_len;
        /*
         * The full-tensor Q4 ID path is the closest arithmetic analogue to IQ2,
         * but PRO Q4 routed tensors are multi-GiB.  Keep it opt-in: even with
         * per-layer command boundaries it is much slower than binding only the
         * six active experts on current M3 Ultra Metal.
         */
        const bool enable_q4_exact_tensor_id =
            getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL &&
            getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL;
        const bool use_q4_exact_tensor_id =
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            !use_q4_group8_experts &&
            !use_q4_group24_experts &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == 384 &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            can_wrap_q4_exact_tensors &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            enable_q4_exact_tensor_id;
        const bool q4_exact_boundary =
            use_q4_exact_tensor_id &&
            g_batch_cb != nil &&
            getenv("DS4_METAL_DISABLE_Q4_EXACT_BOUNDARY") == NULL;
        const bool q4_expert_table_auto =
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
                                                     n_expert,
                                                     gate_tensor_bytes,
                                                     down_tensor_bytes);
        const bool q4_expert_address_auto =
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert,
                                                       n_expert,
                                                       gate_tensor_bytes,
                                                       down_tensor_bytes);
        const bool q4_table_queue_residency =
            ds4_gpu_q4_table_queue_residency_enabled(q4_expert_table_auto ||
                                                     q4_expert_address_auto);
        const bool use_q4_expert_address_table =
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            !use_q4_group8_experts &&
            !use_q4_group24_experts &&
            !use_q4_exact_tensor_id &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == 384 &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil &&
            (getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL ||
             q4_expert_address_auto) &&
            (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL ||
             getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL ||
             q4_table_queue_residency ||
             ds4_gpu_q4_table_model_residency_enabled() ||
             getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) &&
            getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL;
        const bool enable_q4_expert_table =
            getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
            q4_expert_table_auto;
        const bool use_q4_expert_table =
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            !use_q4_group8_experts &&
            !use_q4_group24_experts &&
            !use_q4_exact_tensor_id &&
            !use_q4_expert_address_table &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == 384 &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_table_q4_k_sum6_pipeline != nil &&
            g_moe_table_q4_pair_gate_encoder != nil &&
            g_moe_table_q4_pair_up_encoder != nil &&
            g_moe_table_q4_sum_down_encoder != nil &&
            enable_q4_expert_table &&
            (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL ||
             getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL ||
             q4_table_queue_residency ||
             ds4_gpu_q4_table_model_residency_enabled() ||
             getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) &&
            getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL;
        const bool q4_table_boundary =
            use_q4_expert_table &&
            g_batch_cb != nil &&
            getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL &&
            !q4_table_queue_residency &&
            getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL;
        const bool enable_q4_gather_slots =
            getenv("DS4_METAL_ENABLE_Q4_GATHER_SLOTS") != NULL &&
            getenv("DS4_METAL_DISABLE_Q4_GATHER_SLOTS") == NULL;
        const bool use_q4_gather_slots =
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            !use_q4_group8_experts &&
            !use_q4_group24_experts &&
            !use_q4_exact_tensor_id &&
            !use_q4_expert_address_table &&
            !use_q4_expert_table &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert == q4_group6_expert_group_size * 6u &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_q4_gather_slots6_pipeline != nil &&
            g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil &&
            enable_q4_gather_slots;
        const bool use_q4_selected_slots =
            !force_resident &&
            !use_q4_grouped_experts &&
            !use_q4_group6_experts &&
            !use_q4_group8_experts &&
            !use_q4_group24_experts &&
            !use_q4_exact_tensor_id &&
            !use_q4_expert_address_table &&
            !use_q4_expert_table &&
            !use_q4_gather_slots &&
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            n_total_expert >= 128 &&
            (g_ssd_streaming_mode ||
             (gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
              down_tensor_bytes >= q4_selected_min_tensor_bytes)) &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            ds4_gpu_q4_selected_paths_allowed() &&
            g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil &&
            getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL;
        const bool use_iq2_selected_slots =
            !force_resident &&
            g_ssd_streaming_mode &&
            gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
            down_type == DS4_METAL_TENSOR_Q2_K &&
            n_expert == 6 &&
            n_tokens == 1 &&
            fuse_pair_swiglu &&
            direct_down_sum &&
            g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_slots6_q2_k_sum6_pipeline != nil &&
            getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL;
        const bool use_iq2_stream_addr_table =
            !force_resident &&
            g_ssd_streaming_mode &&
            gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
            down_type == DS4_METAL_TENSOR_IQ2_XXS &&
            n_expert <= DS4_METAL_MAX_ROUTED_EXPERT_USED &&
            n_tokens == 1 &&
            fuse_pair_swiglu &&
            g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_addr_iq2_xxs_pipeline != nil &&
            getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL;
        const bool use_selected_slots =
            use_q4_selected_slots || use_iq2_selected_slots || use_iq2_stream_addr_table;
        id<MTLComputePipelineState> slots_pair_swiglu_pipeline =
            use_iq2_selected_slots ? g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline :
            g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline;
        id<MTLComputePipelineState> slots_sum6_pipeline =
            use_iq2_selected_slots ? g_moe_mul_mv_slots6_q2_k_sum6_pipeline :
            g_moe_mul_mv_slots6_q4_k_sum6_pipeline;
        const char *selected_profile_env = getenv("DS4_METAL_SELECTED_PROFILE");
        if (!selected_profile_env) {
            selected_profile_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE");
        }
        const char *selected_profile_layer_env = getenv("DS4_METAL_SELECTED_PROFILE_LAYER");
        if (!selected_profile_layer_env) {
            selected_profile_layer_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE_LAYER");
        }
        bool selected_profile_layer_match = true;
        if (selected_profile_layer_env && selected_profile_layer_env[0]) {
            char *end = NULL;
            long layer = strtol(selected_profile_layer_env, &end, 10);
            selected_profile_layer_match =
                end && *end == '\0' && layer >= 0 && (uint32_t)layer == layer_index;
        }
        const bool selected_profile =
            use_selected_slots &&
            selected_profile_env != NULL &&
            selected_profile_layer_match;
        const bool q4_selected_shared_event =
            use_q4_selected_slots &&
            getenv("DS4_METAL_Q4_SELECTED_SHARED_EVENT") != NULL;
        const bool q4_selected_base_views =
            use_q4_selected_slots &&
            getenv("DS4_METAL_Q4_SELECTED_USE_BASE_VIEWS") != NULL &&
            getenv("DS4_METAL_Q4_SELECTED_EXACT_VIEWS") == NULL;
        const bool q4_selected_transient_views =
            use_q4_selected_slots &&
            !q4_selected_base_views &&
            getenv("DS4_METAL_Q4_SELECTED_TRANSIENT_VIEWS") != NULL;
        const char *q4_selected_view_mode =
            q4_selected_base_views ? "base" :
            (q4_selected_transient_views ? "transient" : "cached");
        if (!use_selected_slots) {
            g_routed_moe_selected_override_n = 0;
        }
        if (use_q4_expert_address_table) {
            gate_table = ds4_gpu_q4_expert_address_table(model_map,
                                                         model_size,
                                                         gate_offset,
                                                         gate_expert_bytes,
                                                         n_total_expert);
            up_table = ds4_gpu_q4_expert_address_table(model_map,
                                                       model_size,
                                                       up_offset,
                                                       gate_expert_bytes,
                                                       n_total_expert);
            down_table = ds4_gpu_q4_expert_address_table(model_map,
                                                         model_size,
                                                         down_offset,
                                                         down_expert_bytes,
                                                         n_total_expert);
            if (!gate_table || !up_table || !down_table) {
                if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32546);
                return 0;
            }
            q4_table_layer_residency =
                ds4_gpu_q4_expert_layer_residency_set(gate_table,
                                                       up_table,
                                                       down_table,
                                                       q4_expert_address_auto);
        } else if (use_q4_expert_table) {
            gate_table = ds4_gpu_q4_expert_table(model_map,
                                                 model_size,
                                                 gate_offset,
                                                 gate_expert_bytes,
                                                 n_total_expert,
                                                 g_moe_table_q4_pair_gate_encoder);
            up_table = ds4_gpu_q4_expert_table(model_map,
                                               model_size,
                                               up_offset,
                                               gate_expert_bytes,
                                               n_total_expert,
                                               g_moe_table_q4_pair_up_encoder);
            down_table = ds4_gpu_q4_expert_table(model_map,
                                                 model_size,
                                                 down_offset,
                                                 down_expert_bytes,
                                                 n_total_expert,
                                                 g_moe_table_q4_sum_down_encoder);
            if (!gate_table || !up_table || !down_table) {
                if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32573);
                return 0;
            }
            q4_table_layer_residency =
                ds4_gpu_q4_expert_layer_residency_set(gate_table,
                                                       up_table,
                                                       down_table,
                                                       q4_expert_table_auto);
        } else if (use_q4_exact_tensor_id) {
            gate_buf = ds4_gpu_wrap_model_exact_range(model_map,
                                                      model_size,
                                                      gate_offset,
                                                      gate_tensor_bytes,
                                                      &gate_inner);
            up_buf = ds4_gpu_wrap_model_exact_range(model_map,
                                                    model_size,
                                                    up_offset,
                                                    gate_tensor_bytes,
                                                    &up_inner);
            down_buf = ds4_gpu_wrap_model_exact_range(model_map,
                                                      model_size,
                                                      down_offset,
                                                      down_tensor_bytes,
                                                      &down_inner);
            if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32601); return 0; }
        } else if (use_q4_group6_experts || use_q4_gather_slots) {
            if ((uint64_t)q4_group6_expert_group_size > UINT64_MAX / gate_expert_bytes ||
                (uint64_t)q4_group6_expert_group_size > UINT64_MAX / down_expert_bytes) {
                fprintf(stderr, "ds4: Metal routed MoE Q4 group6 byte size overflow\n");
                return 0;
            }
            const uint64_t gate_group_bytes =
                (uint64_t)q4_group6_expert_group_size * gate_expert_bytes;
            const uint64_t down_group_bytes =
                (uint64_t)q4_group6_expert_group_size * down_expert_bytes;
            for (uint32_t i = 0; i < 6; i++) {
                const uint64_t gate_rel = (uint64_t)i * gate_group_bytes;
                const uint64_t down_rel = (uint64_t)i * down_group_bytes;
                if (gate_rel > gate_tensor_bytes ||
                    gate_group_bytes > gate_tensor_bytes - gate_rel ||
                    down_rel > down_tensor_bytes ||
                    down_group_bytes > down_tensor_bytes - down_rel ||
                    gate_rel > UINT64_MAX - gate_offset ||
                    gate_rel > UINT64_MAX - up_offset ||
                    down_rel > UINT64_MAX - down_offset) {
                    fprintf(stderr, "ds4: Metal routed MoE Q4 group6 offset overflow\n");
                    return 0;
                }

                uint64_t group_inner = 0;
                gate_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                gate_offset + gate_rel,
                                                                gate_group_bytes,
                                                                &group_inner);
                gate_group6_offsets[i] = (NSUInteger)group_inner;
                group_inner = 0;
                up_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map,
                                                              model_size,
                                                              up_offset + gate_rel,
                                                              gate_group_bytes,
                                                              &group_inner);
                up_group6_offsets[i] = (NSUInteger)group_inner;
                group_inner = 0;
                down_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                down_offset + down_rel,
                                                                down_group_bytes,
                                                                &group_inner);
                down_group6_offsets[i] = (NSUInteger)group_inner;
                if (!gate_group6_bufs[i] || !up_group6_bufs[i] || !down_group6_bufs[i]) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32643);
                    return 0;
                }
            }
            if (use_q4_gather_slots) {
                if (gate_expert_bytes > NSUIntegerMax ||
                    down_expert_bytes > NSUIntegerMax ||
                    gate_expert_bytes > UINT64_MAX / 6u ||
                    down_expert_bytes > UINT64_MAX / 6u ||
                    6ull * gate_expert_bytes > NSUIntegerMax ||
                    6ull * down_expert_bytes > NSUIntegerMax) {
                    fprintf(stderr, "ds4: Metal routed MoE Q4 gather scratch byte size overflow\n");
                    return 0;
                }
                const NSUInteger gate_slots_bytes = (NSUInteger)(6ull * gate_expert_bytes);
                const NSUInteger down_slots_bytes = (NSUInteger)(6ull * down_expert_bytes);
                if (!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_gate_slots_buffer,
                                                    &g_moe_q4_gate_slots_bytes,
                                                    gate_slots_bytes,
                                                    "ds4_moe_q4_gate_slots") ||
                    !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_up_slots_buffer,
                                                    &g_moe_q4_up_slots_bytes,
                                                    gate_slots_bytes,
                                                    "ds4_moe_q4_up_slots") ||
                    !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_down_slots_buffer,
                                                    &g_moe_q4_down_slots_bytes,
                                                    down_slots_bytes,
                                                    "ds4_moe_q4_down_slots")) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32670);
                    return 0;
                }
                for (uint32_t i = 0; i < 6; i++) {
                    gate_slot_bufs[i] = g_moe_q4_gate_slots_buffer;
                    up_slot_bufs[i] = g_moe_q4_up_slots_buffer;
                    down_slot_bufs[i] = g_moe_q4_down_slots_buffer;
                    gate_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes);
                    up_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes);
                    down_slot_offsets[i] = (NSUInteger)((uint64_t)i * down_expert_bytes);
                }
            }
        } else if (use_q4_group8_experts) {
            if ((uint64_t)q4_group8_expert_group_size > UINT64_MAX / gate_expert_bytes ||
                (uint64_t)q4_group8_expert_group_size > UINT64_MAX / down_expert_bytes) {
                fprintf(stderr, "ds4: Metal routed MoE Q4 group8 byte size overflow\n");
                return 0;
            }
            const uint64_t gate_group_bytes =
                (uint64_t)q4_group8_expert_group_size * gate_expert_bytes;
            const uint64_t down_group_bytes =
                (uint64_t)q4_group8_expert_group_size * down_expert_bytes;
            for (uint32_t i = 0; i < 8; i++) {
                const uint64_t gate_rel = (uint64_t)i * gate_group_bytes;
                const uint64_t down_rel = (uint64_t)i * down_group_bytes;
                if (gate_rel > gate_tensor_bytes ||
                    gate_group_bytes > gate_tensor_bytes - gate_rel ||
                    down_rel > down_tensor_bytes ||
                    down_group_bytes > down_tensor_bytes - down_rel ||
                    gate_rel > UINT64_MAX - gate_offset ||
                    gate_rel > UINT64_MAX - up_offset ||
                    down_rel > UINT64_MAX - down_offset) {
                    fprintf(stderr, "ds4: Metal routed MoE Q4 group8 offset overflow\n");
                    return 0;
                }

                uint64_t group_inner = 0;
                gate_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                gate_offset + gate_rel,
                                                                gate_group_bytes,
                                                                &group_inner);
                gate_group8_offsets[i] = (NSUInteger)group_inner;
                group_inner = 0;
                up_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map,
                                                              model_size,
                                                              up_offset + gate_rel,
                                                              gate_group_bytes,
                                                              &group_inner);
                up_group8_offsets[i] = (NSUInteger)group_inner;
                group_inner = 0;
                down_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map,
                                                                model_size,
                                                                down_offset + down_rel,
                                                                down_group_bytes,
                                                                &group_inner);
                down_group8_offsets[i] = (NSUInteger)group_inner;
                if (!gate_group8_bufs[i] || !up_group8_bufs[i] || !down_group8_bufs[i]) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32727);
                    return 0;
                }
            }
        } else if (use_q4_group24_experts) {
            if ((uint64_t)q4_group24_expert_group_size > UINT64_MAX / gate_expert_bytes ||
                (uint64_t)q4_group24_expert_group_size > UINT64_MAX / down_expert_bytes) {
                fprintf(stderr, "ds4: Metal routed MoE Q4 group24 byte size overflow\n");
                return 0;
            }
            const uint64_t gate_group_bytes =
                (uint64_t)q4_group24_expert_group_size * gate_expert_bytes;
            const uint64_t down_group_bytes =
                (uint64_t)q4_group24_expert_group_size * down_expert_bytes;
            for (uint32_t i = 0; i < 24; i++) {
                const uint64_t gate_rel = (uint64_t)i * gate_group_bytes;
                const uint64_t down_rel = (uint64_t)i * down_group_bytes;
                if (gate_rel > gate_tensor_bytes ||
                    gate_group_bytes > gate_tensor_bytes - gate_rel ||
                    down_rel > down_tensor_bytes ||
                    down_group_bytes > down_tensor_bytes - down_rel ||
                    gate_rel > UINT64_MAX - gate_offset ||
                    gate_rel > UINT64_MAX - up_offset ||
                    down_rel > UINT64_MAX - down_offset) {
                    fprintf(stderr, "ds4: Metal routed MoE Q4 group24 offset overflow\n");
                    return 0;
                }

                uint64_t group_inner = 0;
                gate_group24_bufs[i] = q4_group24_exact_views ?
                    ds4_gpu_wrap_model_exact_range(model_map,
                                                   model_size,
                                                   gate_offset + gate_rel,
                                                   gate_group_bytes,
                                                   &group_inner) :
                    ds4_gpu_wrap_model_range(model_map,
                                             model_size,
                                             gate_offset + gate_rel,
                                             gate_group_bytes,
                                             &group_inner);
                gate_group24_offsets[i] = (NSUInteger)group_inner;
                group_inner = 0;
                up_group24_bufs[i] = q4_group24_exact_views ?
                    ds4_gpu_wrap_model_exact_range(model_map,
                                                   model_size,
                                                   up_offset + gate_rel,
                                                   gate_group_bytes,
                                                   &group_inner) :
                    ds4_gpu_wrap_model_range(model_map,
                                             model_size,
                                             up_offset + gate_rel,
                                             gate_group_bytes,
                                             &group_inner);
                up_group24_offsets[i] = (NSUInteger)group_inner;
                group_inner = 0;
                down_group24_bufs[i] = q4_group24_exact_views ?
                    ds4_gpu_wrap_model_exact_range(model_map,
                                                   model_size,
                                                   down_offset + down_rel,
                                                   down_group_bytes,
                                                   &group_inner) :
                    ds4_gpu_wrap_model_range(model_map,
                                             model_size,
                                             down_offset + down_rel,
                                             down_group_bytes,
                                             &group_inner);
                down_group24_offsets[i] = (NSUInteger)group_inner;
                if (!gate_group24_bufs[i] || !up_group24_bufs[i] || !down_group24_bufs[i]) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32794);
                    return 0;
                }
            }
        } else if (use_selected_slots) {
            const bool selected_timing =
                selected_profile ||
                ds4_gpu_stream_expert_timing_summary_enabled();
            double selected_t0 = selected_timing ? ds4_gpu_now_ms() : 0.0;
            double selected_read_ms = 0.0;
            double selected_sync_ms = 0.0;
            double selected_copy_ms = 0.0;
            double selected_wrap_ms = 0.0;
            uint64_t selected_cache_hits0 = g_stream_expert_cache_hits;
            uint64_t selected_cache_misses0 = g_stream_expert_cache_misses;
            uint64_t selected_cache_wraps0 = g_stream_expert_cache_wraps;
            uint64_t selected_cache_evictions0 = g_stream_expert_cache_evictions;
            const char *selected_id_source = "readback";
            bool selected_ids_available = true;
            bool selected_exec_ids_from_host = false;
            const int stream_expert_cache_size_known =
                ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
                                                             down_expert_bytes);
            const bool use_iq2_full_expert_addr_table =
                use_iq2_selected_slots &&
                ds4_gpu_stream_full_expert_addr_table_requested() &&
                g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
                g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil;
            use_stream_expert_cache =
                !use_iq2_full_expert_addr_table &&
                (use_iq2_selected_slots || use_iq2_stream_addr_table || use_q4_selected_slots) &&
                stream_expert_cache_size_known &&
                ds4_gpu_stream_expert_cache_effective_cap(layer_index,
                                                          n_total_expert,
                                                          n_expert) != 0;
            if (use_iq2_stream_addr_table && !use_stream_expert_cache) {
                fprintf(stderr,
                        "ds4: Metal IQ2/IQ2 streaming decode requires a non-empty expert cache\n");
                if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32831);
                return 0;
            }
            const bool stream_split_ready =
                use_stream_expert_cache &&
                ds4_gpu_stream_expert_split_ready();
            const bool use_stream_compact_addr =
                use_stream_expert_cache &&
                use_iq2_selected_slots &&
                ds4_gpu_stream_compact_addr_requested() &&
                !stream_split_ready &&
                !ds4_gpu_stream_expert_masked_addr_requested() &&
                g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
                g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil;
            use_stream_expert_split_candidate =
                use_stream_expert_cache &&
                use_iq2_selected_slots &&
                !use_stream_compact_addr &&
                stream_split_ready &&
                g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil &&
                g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil;
            const bool use_stream_hit_validator =
                use_stream_expert_cache &&
                use_iq2_selected_slots &&
                ds4_gpu_stream_expert_hit_validator_requested() &&
                g_moe_stream_expert_cache_validate_pipeline != nil &&
                g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
                g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil &&
                ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
                                                         &stream_gate_addr_buf,
                                                         &stream_up_addr_buf,
                                                         &stream_down_addr_buf);
            if (use_iq2_full_expert_addr_table) {
                selected_id_source = "gpu-full-addr";
                selected_ids_available = false;
                g_routed_moe_selected_override_n = 0;
                ds4_gpu_stream_expert_cache_entry *full_entry = NULL;
                if (!ds4_gpu_stream_full_expert_addr_table_prepare(model_map,
                                                                   model_size,
                                                                   layer_index,
                                                                   n_total_expert,
                                                                   gate_offset,
                                                                   up_offset,
                                                                   down_offset,
                                                                   gate_expert_bytes,
                                                                   down_expert_bytes,
                                                                   &stream_gate_addr_buf,
                                                                   &stream_up_addr_buf,
                                                                   &stream_down_addr_buf,
                                                                   &full_entry)) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32880);
                    return 0;
                }
                for (uint32_t i = 0; i < n_expert; i++) {
                    stream_slot_entries[i] = full_entry;
                }
                use_stream_expert_addr_table = true;
            } else {
                const int replayed_selected_ids =
                    ds4_gpu_moe_selected_trace_replay(selected_ids, n_expert);
                if (replayed_selected_ids < 0) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32890);
                    return 0;
                } else if (replayed_selected_ids > 0) {
                    selected_id_source = "replay";
                    selected_exec_ids_from_host = true;
                    g_routed_moe_selected_override_n = 0;
                } else if (g_routed_moe_selected_override_n == n_expert) {
                    memcpy(selected_ids,
                           g_routed_moe_selected_override,
                           (size_t)n_expert * sizeof(selected_ids[0]));
                    selected_id_source = "override";
                    selected_exec_ids_from_host = true;
                    g_routed_moe_selected_override_n = 0;
                } else if (use_stream_hit_validator) {
                    g_routed_moe_selected_override_n = 0;
                    uint32_t validator_all_cached = 0;
                    uint32_t validator_miss_mask = 0;
                    uint32_t validator_invalid_mask = 0;
                    if (!ds4_gpu_stream_expert_cache_validate_selected(selected,
                                                                       stream_gate_addr_buf,
                                                                       stream_up_addr_buf,
                                                                       stream_down_addr_buf,
                                                                       n_total_expert,
                                                                       n_expert,
                                                                       selected_ids,
                                                                       &validator_all_cached,
                                                                       &validator_miss_mask,
                                                                       &validator_invalid_mask)) {
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32917);
                        return 0;
                    }
                    selected_id_source =
                        validator_invalid_mask != 0 ? "validator-invalid" :
                        (validator_miss_mask != 0 ? "validator-miss" :
                        (validator_all_cached != 0 ? "validator-hit" : "validator-miss"));
                } else {
                    g_routed_moe_selected_override_n = 0;
                    if (g_batch_cb != nil) {
                        double selected_boundary_t0 =
                            selected_timing ? ds4_gpu_now_ms() : 0.0;
                        if (q4_selected_shared_event) {
                            if (ds4_gpu_signal_batch_and_wait_event("selected-id readback") == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32942); return 0; }
                        } else if (ds4_gpu_end_commands() == 0) {
                            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32931);
                            return 0;
                        }
                        if (selected_timing) {
                            selected_sync_ms +=
                                ds4_gpu_now_ms() - selected_boundary_t0;
                        }
                        double selected_copy_t0 =
                            selected_timing ? ds4_gpu_now_ms() : 0.0;
                        if (ds4_gpu_tensor_read(selected,
                                                0,
                                                selected_ids,
                                                (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) {
                            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32943);
                            return 0;
                        }
                        if (selected_timing) {
                            selected_copy_ms +=
                                ds4_gpu_now_ms() - selected_copy_t0;
                        }
                        if (!q4_selected_shared_event) {
                            selected_boundary_t0 =
                                selected_timing ? ds4_gpu_now_ms() : 0.0;
                            if (ds4_gpu_begin_commands() == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32967); return 0; }
                            if (selected_timing) {
                                selected_sync_ms +=
                                    ds4_gpu_now_ms() - selected_boundary_t0;
                            }
                        }
                    } else {
                        double selected_copy_t0 =
                            selected_timing ? ds4_gpu_now_ms() : 0.0;
                        if (ds4_gpu_tensor_read(selected,
                                                0,
                                                selected_ids,
                                                (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) {
                            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32965);
                            return 0;
                        }
                        if (selected_timing) {
                            selected_copy_ms +=
                                ds4_gpu_now_ms() - selected_copy_t0;
                        }
                    }
                }
            }
            if (selected_timing) {
                selected_read_ms = selected_sync_ms + selected_copy_ms;
                selected_t0 = ds4_gpu_now_ms();
            }

            if (selected_ids_available) {
                for (uint32_t i = 0; i < n_expert; i++) {
                    if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
                        fprintf(stderr,
                                "ds4: Metal routed MoE selected expert id %d is outside 0..%u\n",
                                selected_ids[i],
                                n_total_expert);
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32986);
                        return 0;
                    }
                }
                ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index,
                                                                  selected_ids,
                                                                  n_expert);
                if (!ds4_gpu_moe_selected_trace_record(selected_ids, n_expert) ||
                    !ds4_gpu_moe_selected_hotlist_record(layer_index,
                                                         selected_ids,
                                                         n_expert,
                                                         n_total_expert)) {
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32997);
                    return 0;
                }

                for (uint32_t i = 0; i < n_expert; i++) {
                    const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i];
                    if (expert_id > UINT64_MAX / gate_expert_bytes ||
                        expert_id > UINT64_MAX / down_expert_bytes) {
                        fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n");
                        return 0;
                    }
                    const uint64_t gate_rel = expert_id * gate_expert_bytes;
                    const uint64_t down_rel = expert_id * down_expert_bytes;
                    if (gate_rel > UINT64_MAX - gate_offset ||
                        gate_rel > UINT64_MAX - up_offset ||
                        down_rel > UINT64_MAX - down_offset) {
                        fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n");
                        return 0;
                    }
                    stream_gate_abs_offsets[i] = gate_offset + gate_rel;
                    stream_up_abs_offsets[i] = up_offset + gate_rel;
                    stream_down_abs_offsets[i] = down_offset + down_rel;

                    if (use_stream_expert_cache) {
                        ds4_gpu_stream_expert_cache_entry *entry = NULL;
                        entry = ds4_gpu_stream_expert_cache_peek(model_map,
                                                                 model_size,
                                                                 layer_index,
                                                                 (uint32_t)selected_ids[i],
                                                                 n_total_expert,
                                                                 n_expert,
                                                                 stream_gate_abs_offsets[i],
                                                                 stream_up_abs_offsets[i],
                                                                 stream_down_abs_offsets[i],
                                                                 gate_expert_bytes,
                                                                 down_expert_bytes);
                        if (!entry) {
                            stream_expert_missing_mask |= 1u << i;
                            continue;
                        }
                        stream_expert_resident_mask |= 1u << i;
                        stream_slot_entries[i] = entry;
                        gate_slot_bufs[i] = entry->gate_buffer;
                        gate_slot_offsets[i] = entry->gate_inner;
                        up_slot_bufs[i] = entry->up_buffer;
                        up_slot_offsets[i] = entry->up_inner;
                        down_slot_bufs[i] = entry->down_buffer;
                        down_slot_offsets[i] = entry->down_inner;
                        continue;
                    }

                    uint64_t slot_inner = 0;
                    gate_slot_bufs[i] = q4_selected_base_views ?
                        ds4_gpu_wrap_model_range(model_map,
                                                 model_size,
                                                 gate_offset + gate_rel,
                                                 gate_expert_bytes,
                                                 &slot_inner) :
                        (q4_selected_transient_views ?
                            ds4_gpu_wrap_model_exact_range_transient(model_map,
                                                                     model_size,
                                                                     gate_offset + gate_rel,
                                                                     gate_expert_bytes,
                                                                     &slot_inner) :
                            ds4_gpu_wrap_model_exact_range(model_map,
                                                           model_size,
                                                           gate_offset + gate_rel,
                                                           gate_expert_bytes,
                                                           &slot_inner));
                    gate_slot_offsets[i] = (NSUInteger)slot_inner;
                    slot_inner = 0;
                    up_slot_bufs[i] = q4_selected_base_views ?
                        ds4_gpu_wrap_model_range(model_map,
                                                 model_size,
                                                 up_offset + gate_rel,
                                                 gate_expert_bytes,
                                                 &slot_inner) :
                        (q4_selected_transient_views ?
                            ds4_gpu_wrap_model_exact_range_transient(model_map,
                                                                     model_size,
                                                                     up_offset + gate_rel,
                                                                     gate_expert_bytes,
                                                                     &slot_inner) :
                            ds4_gpu_wrap_model_exact_range(model_map,
                                                           model_size,
                                                           up_offset + gate_rel,
                                                           gate_expert_bytes,
                                                           &slot_inner));
                    up_slot_offsets[i] = (NSUInteger)slot_inner;
                    slot_inner = 0;
                    down_slot_bufs[i] = q4_selected_base_views ?
                        ds4_gpu_wrap_model_range(model_map,
                                                 model_size,
                                                 down_offset + down_rel,
                                                 down_expert_bytes,
                                                 &slot_inner) :
                        (q4_selected_transient_views ?
                            ds4_gpu_wrap_model_exact_range_transient(model_map,
                                                                     model_size,
                                                                     down_offset + down_rel,
                                                                     down_expert_bytes,
                                                                     &slot_inner) :
                            ds4_gpu_wrap_model_exact_range(model_map,
                                                           model_size,
                                                           down_offset + down_rel,
                                                           down_expert_bytes,
                                                           &slot_inner));
                    down_slot_offsets[i] = (NSUInteger)slot_inner;
                    if (!gate_slot_bufs[i] || !up_slot_bufs[i] || !down_slot_bufs[i]) {
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33105);
                        return 0;
                    }
                }
            }
            if (use_stream_expert_cache) {
                use_stream_expert_addr_table =
                    ((use_iq2_selected_slots &&
                      ds4_gpu_stream_expert_addr_table_kernel_requested() &&
                      g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
                      g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil) ||
                     use_iq2_stream_addr_table) &&
                    ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
                                                             &stream_gate_addr_buf,
                                                             &stream_up_addr_buf,
                                                             &stream_down_addr_buf);
                if (use_iq2_stream_addr_table && !use_stream_expert_addr_table) {
                    fprintf(stderr,
                            "ds4: Metal IQ2/IQ2 streaming decode could not prepare expert address buffers\n");
                    if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33123);
                    return 0;
                }
                use_stream_expert_masked_addr_table =
                    use_stream_expert_addr_table &&
                    use_iq2_selected_slots &&
                    ds4_gpu_stream_expert_masked_addr_requested() &&
                    g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil &&
                    g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil;
                use_stream_expert_split_deferred =
                    use_stream_expert_split_candidate &&
                    use_stream_expert_masked_addr_table &&
                    stream_expert_resident_mask != 0 &&
                    stream_expert_missing_mask != 0 &&
                    ds4_gpu_stream_expert_split_worthwhile(stream_expert_resident_mask,
                                                           stream_expert_missing_mask) &&
                    g_batch_cb != nil &&
                    getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL;
                if (use_stream_expert_split_deferred) {
                    const ds4_gpu_stream_expert_table table = {
                        .model_map = model_map,
                        .model_size = model_size,
                        .layer = layer_index,
                        .n_total_expert = n_total_expert,
                        .gate_offset = gate_offset,
                        .up_offset = up_offset,
                        .down_offset = down_offset,
                        .gate_expert_bytes = gate_expert_bytes,
                        .down_expert_bytes = down_expert_bytes,
                    };
                    if (!ds4_gpu_stream_expert_cache_begin_selected_load(
                                &table,
                                selected_ids,
                                n_expert)) {
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33156);
                        return 0;
                    }
                }
                if (stream_expert_missing_mask != 0 &&
                    !use_stream_expert_split_deferred) {
                    if (!ds4_gpu_stream_expert_cache_load_selected_missing(
                            model_map,
                            model_size,
                            layer_index,
                            selected_ids,
                            n_total_expert,
                            n_expert,
                            stream_gate_abs_offsets,
                            stream_up_abs_offsets,
                            stream_down_abs_offsets,
                            gate_expert_bytes,
                            down_expert_bytes,
                            stream_expert_missing_mask,
                            stream_slot_entries)) {
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33175);
                        return 0;
                    }
                    for (uint32_t i = 0; i < n_expert; i++) {
                        if ((stream_expert_missing_mask & (1u << i)) == 0) continue;
                        ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i];
                        if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); return 0; }
                        gate_slot_bufs[i] = entry->gate_buffer;
                        gate_slot_offsets[i] = entry->gate_inner;
                        up_slot_bufs[i] = entry->up_buffer;
                        up_slot_offsets[i] = entry->up_inner;
                        down_slot_bufs[i] = entry->down_buffer;
                        down_slot_offsets[i] = entry->down_inner;
                    }
                }
                if (use_iq2_stream_addr_table && use_stream_expert_addr_table) {
                    for (uint32_t i = 0; i < n_expert; i++) {
                        ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i];
                        if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33214); return 0; }
                        if (!ds4_gpu_stream_expert_cache_set_addr_slot_raw(
                                    layer_index,
                                    (uint32_t)selected_ids[i],
                                    entry->gate_buffer,
                                    entry->gate_inner,
                                    entry->up_buffer,
                                    entry->up_inner,
                                    entry->down_buffer,
                                    entry->down_inner)) {
                            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202);
                            return 0;
                        }
                    }
                }
                ds4_gpu_stream_expert_cache_prune_layer(layer_index,
                                                        n_total_expert,
                                                        n_expert,
                                                        selected_ids,
                                                        n_expert);
                ds4_gpu_stream_expert_cache_prune_global(layer_index,
                                                         selected_ids,
                                                         n_expert);
                if (use_stream_compact_addr) {
                    id<MTLBuffer> compact_selected = nil;
                    if (!ds4_gpu_stream_compact_addr_prepare(layer_index,
                                                             stream_slot_entries,
                                                             n_expert,
                                                             &stream_gate_addr_buf,
                                                             &stream_up_addr_buf,
                                                             &stream_down_addr_buf,
                                                             &compact_selected)) {
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33223);
                        return 0;
                    }
                    selected_exec_buf = compact_selected;
                    selected_exec_off = 0;
                    use_stream_expert_addr_table = true;
                    use_stream_expert_masked_addr_table = false;
                    use_stream_compact_addr_table = true;
                }
                if (use_stream_expert_addr_table &&
                    !use_stream_compact_addr_table &&
                    selected_exec_ids_from_host) {
                    if (!ds4_gpu_stream_selected_ids_prepare(layer_index,
                                                             selected_ids,
                                                             n_expert,
                                                             &selected_exec_buf,
                                                             &selected_exec_off)) {
                        if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33239);
                        return 0;
                    }
                }
            }
            if (selected_timing) {
                selected_wrap_ms = ds4_gpu_now_ms() - selected_t0;
                if (use_stream_expert_cache) {
                    ds4_gpu_stream_expert_timing_note_cache_class(
                            stream_expert_resident_mask,
                            stream_expert_missing_mask);
                }
                ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms,
                                                           selected_copy_ms,
                                                           selected_wrap_ms);
            }
            if (selected_profile) {
                const uint64_t selected_cache_hits =
                    g_stream_expert_cache_hits - selected_cache_hits0;
                const uint64_t selected_cache_misses =
                    g_stream_expert_cache_misses - selected_cache_misses0;
                const uint64_t selected_cache_wraps =
                    g_stream_expert_cache_wraps - selected_cache_wraps0;
                const uint64_t selected_cache_evictions =
                    g_stream_expert_cache_evictions - selected_cache_evictions0;
                const char *selected_path =
                    use_iq2_stream_addr_table ? "iq2/iq2" :
                    (use_iq2_selected_slots ? "iq2/q2" : "q4/q4");
                const char *selected_view_mode =
                    use_stream_expert_split_deferred ? "stream-split" :
                    use_stream_expert_masked_addr_table ? "stream-addr-mask" :
                    use_stream_compact_addr_table ? "stream-compact-addr" :
                    use_stream_expert_addr_table ? "stream-addr" :
                    (use_stream_expert_cache ? "stream-cache" :
                    (use_iq2_selected_slots ? "exact-cache" : q4_selected_view_mode));
                fprintf(stderr,
                        "ds4: Metal selected views layer=%u path=%s mode=%s ids=%s "
                        "experts=%d,%d,%d,%d,%d,%d expert_gate=%.2f MiB "
                        "expert_down=%.2f MiB read=%.3f ms bind=%.3f ms "
                        "cache_hits=%llu cache_misses=%llu cache_wraps=%llu cache_evictions=%llu\n",
                        layer_index,
                        selected_path,
                        selected_view_mode,
                        selected_id_source,
                        selected_ids_available ? selected_ids[0] : -1,
                        selected_ids_available ? selected_ids[1] : -1,
                        selected_ids_available ? selected_ids[2] : -1,
                        selected_ids_available ? selected_ids[3] : -1,
                        selected_ids_available ? selected_ids[4] : -1,
                        selected_ids_available ? selected_ids[5] : -1,
                        ds4_gpu_mib(gate_expert_bytes),
                        ds4_gpu_mib(down_expert_bytes),
                        selected_read_ms,
                        selected_wrap_ms,
                        (unsigned long long)selected_cache_hits,
                        (unsigned long long)selected_cache_misses,
                        (unsigned long long)selected_cache_wraps,
                        (unsigned long long)selected_cache_evictions);
            }
        } else if (!use_q4_grouped_experts) {
            gate_buf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, gate_tensor_bytes, &gate_inner);
            up_buf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, gate_tensor_bytes, &up_inner);
            down_buf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, down_tensor_bytes, &down_inner);
            if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33326); return 0; }
        }
        if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33305);
                return 0;
            }
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33337); return 0; }
        if ((use_q4_expert_address_table || use_q4_expert_table) &&
            !ds4_gpu_use_model_residency_set(cb)) {
            if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33314);
            return 0;
        }
        if (!q4_table_queue_residency &&
            q4_table_layer_residency &&
            [cb respondsToSelector:@selector(useResidencySet:)]) {
            [cb useResidencySet:q4_table_layer_residency];
        }

        const bool moe_one_stage_profile =
            g_batch_cb != nil &&
            ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_ONE_STAGE_PROFILE",
                                                    "DS4_METAL_MOE_ONE_STAGE_PROFILE_LAYER",
                                                    layer_index);
        const char *moe_one_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER");
        const char *moe_one_path =
            use_q4_grouped_experts ? "q4_grouped_pair_swiglu" :
            use_q4_group6_experts ? "q4_group6_pair_swiglu" :
            use_q4_group8_experts ? "q4_group8_pair_swiglu" :
            use_q4_group24_experts ? "q4_group24_split_gate_up" :
            use_q4_exact_tensor_id ? "q4_exact_pair_swiglu" :
            use_q4_expert_address_table ? "q4_addr_pair_swiglu" :
            use_q4_expert_table ? "q4_table_pair_swiglu" :
            use_q4_gather_slots ? "q4_gather_slots6_pair_swiglu" :
            use_stream_expert_split_deferred ? "iq2_stream_split_pair_swiglu" :
            use_stream_expert_masked_addr_table ? "iq2_stream_addr_mask_pair_swiglu" :
            use_stream_expert_addr_table ? "iq2_stream_addr_pair_swiglu" :
            use_iq2_selected_slots ? "iq2_slots6_pair_swiglu" :
            use_q4_selected_slots ? "q4_slots6_pair_swiglu" :
            (fuse_pair_swiglu ? "pair_swiglu" :
            ((!g_quality_mode &&
              ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) ||
               (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline))) ? "pair" : "single"));
        double moe_one_stage_t0 = moe_one_stage_profile ? ds4_gpu_now_ms() : 0.0;
        if (moe_one_stage_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33349);
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33380); return 0; }
            moe_one_stage_t0 = ds4_gpu_now_ms();
        }
#define DS4_METAL_PROFILE_MOE_ONE_STAGE(name) do { \
            if (ok && moe_one_stage_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = 0; \
                } else { \
                    const char *stage_name = (name); \
                    const double now_ms = ds4_gpu_now_ms(); \
                    const int print_stage = \
                        !moe_one_stage_filter || !moe_one_stage_filter[0] || \
                        strstr(stage_name, moe_one_stage_filter) != NULL; \
                    if (print_stage) { \
                        fprintf(stderr, \
                                "ds4: Metal routed MoE one stage layer=%u pairs=%u experts=%u " \
                                "gate=%s down=%s path=%s %s=%.3f ms\n", \
                                layer_index, pair_rows, n_expert, \
                                ds4_gpu_metal_tensor_type_name(gate_type), \
                                ds4_gpu_metal_tensor_type_name(down_type), \
                                moe_one_path, \
                                stage_name, now_ms - moe_one_stage_t0); \
                    } \
                    moe_one_stage_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = 0; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb) ok = 0; \
                    } \
                } \
            } \
        } while (0)
        if (use_q4_gather_slots) {
            ds4_gpu_q4_gather_slots6_args gate_gather_args = {
                .expert_bytes = gate_expert_bytes,
                .group_size = q4_group6_expert_group_size,
                .n_slots = n_expert,
            };
            ds4_gpu_q4_gather_slots6_args down_gather_args = {
                .expert_bytes = down_expert_bytes,
                .group_size = q4_group6_expert_group_size,
                .n_slots = n_expert,
            };
            ok = ds4_gpu_encode_q4_gather_slots6(cb,
                                                  g_moe_q4_gather_slots6_pipeline,
                                                  &gate_gather_args,
                                                  gate_group6_bufs,
                                                  gate_group6_offsets,
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  g_moe_q4_gate_slots_buffer,
                                                  0) &&
                 ds4_gpu_encode_q4_gather_slots6(cb,
                                                  g_moe_q4_gather_slots6_pipeline,
                                                  &gate_gather_args,
                                                  up_group6_bufs,
                                                  up_group6_offsets,
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  g_moe_q4_up_slots_buffer,
                                                  0) &&
                 ds4_gpu_encode_q4_gather_slots6(cb,
                                                  g_moe_q4_gather_slots6_pipeline,
                                                  &down_gather_args,
                                                  down_group6_bufs,
                                                  down_group6_offsets,
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  g_moe_q4_down_slots_buffer,
                                                  0);
        }
        if (use_q4_gather_slots) {
            DS4_METAL_PROFILE_MOE_ONE_STAGE("q4_gather");
        }
        if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33455); return 0; }
        if (use_q4_grouped_experts) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) {
                const uint32_t expert_count =
                    q4_expert_group_size < n_total_expert - expert_base ?
                    q4_expert_group_size : n_total_expert - expert_base;
                if ((uint64_t)expert_base > UINT64_MAX / gate_expert_bytes ||
                    (uint64_t)expert_count > UINT64_MAX / gate_expert_bytes) {
                    ok = 0;
                    break;
                }
                const uint64_t group_rel = (uint64_t)expert_base * gate_expert_bytes;
                const uint64_t group_bytes = (uint64_t)expert_count * gate_expert_bytes;
                if (group_rel > UINT64_MAX - gate_offset ||
                    group_rel > UINT64_MAX - up_offset) {
                    ok = 0;
                    break;
                }
                uint64_t gate_group_inner = 0;
                uint64_t up_group_inner = 0;
                id<MTLBuffer> gate_group_buf =
                    q4_grouped_cache_views ?
                    ds4_gpu_wrap_model_exact_range(model_map,
                                                   model_size,
                                                   gate_offset + group_rel,
                                                   group_bytes,
                                                   &gate_group_inner) :
                    ds4_gpu_wrap_model_exact_range_transient(model_map,
                                                             model_size,
                                                             gate_offset + group_rel,
                                                             group_bytes,
                                                             &gate_group_inner);
                id<MTLBuffer> up_group_buf =
                    q4_grouped_cache_views ?
                    ds4_gpu_wrap_model_exact_range(model_map,
                                                   model_size,
                                                   up_offset + group_rel,
                                                   group_bytes,
                                                   &up_group_inner) :
                    ds4_gpu_wrap_model_exact_range_transient(model_map,
                                                             model_size,
                                                             up_offset + group_rel,
                                                             group_bytes,
                                                             &up_group_inner);
                if (!gate_group_buf || !up_group_buf) {
                    ok = 0;
                    break;
                }
                ds4_gpu_moe_expert_group_args group_args = {
                    .expert_base = expert_base,
                    .expert_count = expert_count,
                    .accumulate = 0,
                    .pad0 = 0,
                };
                ok = ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(cb,
                                                                g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline,
                                                                &gate_args,
                                                                &act_args,
                                                                &group_args,
                                                                gate_group_buf,
                                                                (NSUInteger)gate_group_inner,
                                                                up_group_buf,
                                                                (NSUInteger)up_group_inner,
                                                                xbuf,
                                                                ds4_gpu_tensor_offset(x),
                                                                gatebuf,
                                                                ds4_gpu_tensor_offset(gate),
                                                                upbuf,
                                                                ds4_gpu_tensor_offset(up),
                                                                midbuf,
                                                                ds4_gpu_tensor_offset(mid),
                                                                selectedbuf,
                                                                ds4_gpu_tensor_offset(selected),
                                                                weightsbuf,
                                                                ds4_gpu_tensor_offset(weights),
                                                                gate_smem,
                                                                2,
                                                                false);
            }
        } else if (use_q4_expert_address_table) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(cb,
                                                           g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline,
                                                           &gate_args,
                                                           &act_args,
                                                           gate_table,
                                                           up_table,
                                                           xbuf,
                                                           ds4_gpu_tensor_offset(x),
                                                           gatebuf,
                                                           ds4_gpu_tensor_offset(gate),
                                                           upbuf,
                                                           ds4_gpu_tensor_offset(up),
                                                           midbuf,
                                                           ds4_gpu_tensor_offset(mid),
                                                           selectedbuf,
                                                           ds4_gpu_tensor_offset(selected),
                                                           weightsbuf,
                                                           ds4_gpu_tensor_offset(weights),
                                                           gate_smem,
                                                           2,
                                                           false);
        } else if (use_q4_expert_table) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb,
                                                            g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline,
                                                            &gate_args,
                                                            &act_args,
                                                            gate_table,
                                                            up_table,
                                                            xbuf,
                                                            ds4_gpu_tensor_offset(x),
                                                            gatebuf,
                                                            ds4_gpu_tensor_offset(gate),
                                                            upbuf,
                                                            ds4_gpu_tensor_offset(up),
                                                            midbuf,
                                                            ds4_gpu_tensor_offset(mid),
                                                            selectedbuf,
                                                            ds4_gpu_tensor_offset(selected),
                                                            weightsbuf,
                                                            ds4_gpu_tensor_offset(weights),
                                                            gate_smem,
                                                            2,
                                                            false,
                                                            q4_table_queue_residency);
        } else if (use_q4_group6_experts) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_group6_pair_swiglu(cb,
                                                            g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline,
                                                            &gate_args,
                                                            &act_args,
                                                            gate_group6_bufs,
                                                            gate_group6_offsets,
                                                            up_group6_bufs,
                                                            up_group6_offsets,
                                                            xbuf,
                                                            ds4_gpu_tensor_offset(x),
                                                            gatebuf,
                                                            ds4_gpu_tensor_offset(gate),
                                                            upbuf,
                                                            ds4_gpu_tensor_offset(up),
                                                            midbuf,
                                                            ds4_gpu_tensor_offset(mid),
                                                            selectedbuf,
                                                            ds4_gpu_tensor_offset(selected),
                                                            weightsbuf,
                                                            ds4_gpu_tensor_offset(weights),
                                                            gate_smem,
                                                            2,
                                                            false);
        } else if (use_q4_group8_experts) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_group8_pair_swiglu(cb,
                                                            g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline,
                                                            &gate_args,
                                                            &act_args,
                                                            gate_group8_bufs,
                                                            gate_group8_offsets,
                                                            up_group8_bufs,
                                                            up_group8_offsets,
                                                            xbuf,
                                                            ds4_gpu_tensor_offset(x),
                                                            gatebuf,
                                                            ds4_gpu_tensor_offset(gate),
                                                            upbuf,
                                                            ds4_gpu_tensor_offset(up),
                                                            midbuf,
                                                            ds4_gpu_tensor_offset(mid),
                                                            selectedbuf,
                                                            ds4_gpu_tensor_offset(selected),
                                                            weightsbuf,
                                                            ds4_gpu_tensor_offset(weights),
                                                            gate_smem,
                                                            2,
                                                            false);
        } else if (use_q4_group24_experts) {
            ok = ds4_gpu_encode_mul_mv_group24_id(cb,
                                                  g_moe_mul_mv_group24_q4_k_id_pipeline,
                                                  &gate_args,
                                                  gate_group24_bufs,
                                                  gate_group24_offsets,
                                                  xbuf,
                                                  ds4_gpu_tensor_offset(x),
                                                  gatebuf,
                                                  ds4_gpu_tensor_offset(gate),
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  gate_smem,
                                                  2,
                                                  false) &&
                 ds4_gpu_encode_mul_mv_group24_id(cb,
                                                  g_moe_mul_mv_group24_q4_k_id_pipeline,
                                                  &gate_args,
                                                  up_group24_bufs,
                                                  up_group24_offsets,
                                                  xbuf,
                                                  ds4_gpu_tensor_offset(x),
                                                  upbuf,
                                                  ds4_gpu_tensor_offset(up),
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  gate_smem,
                                                  2,
                                                  false);
        } else if (use_q4_gather_slots || use_selected_slots) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            if (use_stream_expert_addr_table) {
                if (use_stream_expert_masked_addr_table) {
                    if (use_stream_expert_split_deferred) {
                        const bool stream_split_profile =
                            getenv("DS4_METAL_STREAMING_EXPERT_SPLIT_PROFILE") != NULL;
                        const bool stream_split_timing =
                            stream_split_profile ||
                            ds4_gpu_stream_expert_timing_summary_enabled();
                        double stream_split_t0 =
                            stream_split_timing ? ds4_gpu_now_ms() : 0.0;
                        ds4_gpu_stream_expert_split_args resident_pair_args = {
                            .active_mask = stream_expert_resident_mask,
                            .accumulate = 0u,
                        };
                        ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb,
                                                                               g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline,
                                                                               &gate_args,
                                                                               &act_args,
                                                                               &resident_pair_args,
                                                                               stream_slot_entries,
                                                                               stream_gate_addr_buf,
                                                                               stream_up_addr_buf,
                                                                               xbuf,
                                                                               ds4_gpu_tensor_offset(x),
                                                                               gatebuf,
                                                                               ds4_gpu_tensor_offset(gate),
                                                                               upbuf,
                                                                               ds4_gpu_tensor_offset(up),
                                                                               midbuf,
                                                                               ds4_gpu_tensor_offset(mid),
                                                                               selected_exec_buf,
                                                                               selected_exec_off,
                                                                               weightsbuf,
                                                                               ds4_gpu_tensor_offset(weights),
                                                                               gate_smem,
                                                                               2,
                                                                               false);
                        if (ok) {
                            ok = ds4_gpu_flush_commands();
                            if (ok) {
                                cb = ds4_gpu_command_buffer(&owned);
                                if (!cb) ok = 0;
                            }
                        }
                        const double stream_split_resident_ms =
                            stream_split_timing ? ds4_gpu_now_ms() - stream_split_t0 : 0.0;
                        if (stream_split_timing) stream_split_t0 = ds4_gpu_now_ms();
                        const double stream_split_missing_start_ms = stream_split_t0;
                        double stream_split_missing_load_ms = 0.0;
                        double stream_split_missing_slot_ms = 0.0;
                        double stream_split_missing_prune_ms = 0.0;
                        double stream_split_missing_addr_ms = 0.0;
                        double stream_split_missing_wait_ms = 0.0;
                        if (ok) {
                            ok = ds4_gpu_stream_expert_cache_load_selected_missing(
                                    model_map,
                                    model_size,
                                    layer_index,
                                    selected_ids,
                                    n_total_expert,
                                    n_expert,
                                    stream_gate_abs_offsets,
                                    stream_up_abs_offsets,
                                    stream_down_abs_offsets,
                                    gate_expert_bytes,
                                    down_expert_bytes,
                                    stream_expert_missing_mask,
                                    stream_slot_entries);
                            if (stream_split_timing) {
                                const double now_ms = ds4_gpu_now_ms();
                                stream_split_missing_load_ms =
                                    now_ms - stream_split_t0;
                                stream_split_t0 = now_ms;
                            }
                            if (ok) {
                                for (uint32_t i = 0; i < n_expert; i++) {
                                    if ((stream_expert_missing_mask & (1u << i)) == 0) continue;
                                    ds4_gpu_stream_expert_cache_entry *entry =
                                        stream_slot_entries[i];
                                    if (!entry) {
                                        ok = 0;
                                        break;
                                    }
                                    gate_slot_bufs[i] = entry->gate_buffer;
                                    gate_slot_offsets[i] = entry->gate_inner;
                                    up_slot_bufs[i] = entry->up_buffer;
                                    up_slot_offsets[i] = entry->up_inner;
                                    down_slot_bufs[i] = entry->down_buffer;
                                    down_slot_offsets[i] = entry->down_inner;
                                }
                            }
                        }
                        if (stream_split_timing) {
                            const double now_ms = ds4_gpu_now_ms();
                            stream_split_missing_slot_ms =
                                now_ms - stream_split_t0;
                            stream_split_t0 = now_ms;
                        }
                        if (ok) {
                            ds4_gpu_stream_expert_cache_prune_layer(layer_index,
                                                                    n_total_expert,
                                                                    n_expert,
                                                                    selected_ids,
                                                                    n_expert);
                            ds4_gpu_stream_expert_cache_prune_global(layer_index,
                                                                     selected_ids,
                                                                     n_expert);
                            if (stream_split_timing) {
                                const double now_ms = ds4_gpu_now_ms();
                                stream_split_missing_prune_ms =
                                    now_ms - stream_split_t0;
                                stream_split_t0 = now_ms;
                            }
                            ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
                                                                          &stream_gate_addr_buf,
                                                                          &stream_up_addr_buf,
                                                                          &stream_down_addr_buf);
                            if (stream_split_timing) {
                                const double now_ms = ds4_gpu_now_ms();
                                stream_split_missing_addr_ms =
                                    now_ms - stream_split_t0;
                                stream_split_t0 = now_ms;
                            }
                        }
                        if (ok) {
                            /*
                             * The resident stage was submitted before the
                             * CPU read of missing experts so I/O can overlap
                             * with GPU work. The missing stage reuses the same
                             * gate/up/mid scratch buffers, so it must not
                             * execute until the resident command buffer has
                             * finished. The down/sum pass is issued once after
                             * all six mid slots exist; this keeps the final
                             * accumulation order stable regardless of the
                             * resident/missing split.
                             */
                            ok = ds4_gpu_wait_pending_command_buffers(
                                    "streaming expert split resident");
                            if (stream_split_timing) {
                                const double now_ms = ds4_gpu_now_ms();
                                stream_split_missing_wait_ms =
                                    now_ms - stream_split_t0;
                                stream_split_t0 = now_ms;
                            }
                        }
                        const double stream_split_missing_ms =
                            stream_split_timing ?
                                ds4_gpu_now_ms() - stream_split_missing_start_ms :
                                0.0;
                        ds4_gpu_stream_expert_split_args missing_pair_args = {
                            .active_mask = stream_expert_missing_mask,
                            .accumulate = 0u,
                        };
                        ds4_gpu_stream_expert_split_args all_down_args = {
                            .active_mask = 0x3fu,
                            .accumulate = 0u,
                        };
                        if (ok) {
                            ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb,
                                                                                   g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline,
                                                                                   &gate_args,
                                                                                   &act_args,
                                                                                   &missing_pair_args,
                                                                                   stream_slot_entries,
                                                                                   stream_gate_addr_buf,
                                                                                   stream_up_addr_buf,
                                                                                   xbuf,
                                                                                   ds4_gpu_tensor_offset(x),
                                                                                   gatebuf,
                                                                                   ds4_gpu_tensor_offset(gate),
                                                                                   upbuf,
                                                                                   ds4_gpu_tensor_offset(up),
                                                                                   midbuf,
                                                                                   ds4_gpu_tensor_offset(mid),
                                                                                   selected_exec_buf,
                                                                                   selected_exec_off,
                                                                                   weightsbuf,
                                                                                   ds4_gpu_tensor_offset(weights),
                                                                                   gate_smem,
                                                                                   2,
                                                                                   false);
                        }
                        if (ok) {
                            ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb,
                                                                           g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline,
                                                                           &down_args,
                                                                           &all_down_args,
                                                                           stream_slot_entries,
                                                                           stream_down_addr_buf,
                                                                           midbuf,
                                                                           ds4_gpu_tensor_offset(mid),
                                                                           outbuf,
                                                                           ds4_gpu_tensor_offset(out),
                                                                           selected_exec_buf,
                                                                           selected_exec_off,
                                                                           down_smem,
                                                                           2);
                        }
                        stream_expert_split_completed = ok;
                        if (stream_split_timing) {
                            ds4_gpu_stream_expert_timing_note_split(
                                    stream_expert_resident_mask,
                                    stream_expert_missing_mask,
                                    stream_split_resident_ms,
                                    stream_split_missing_ms);
                            ds4_gpu_stream_expert_timing_note_split_missing_detail(
                                    stream_split_missing_load_ms,
                                    stream_split_missing_slot_ms,
                                    stream_split_missing_prune_ms,
                                    stream_split_missing_addr_ms,
                                    stream_split_missing_wait_ms);
                        }
                        if (stream_split_profile) {
                            fprintf(stderr,
                                    "ds4: Metal streaming expert split layer=%u "
                                    "resident=0x%02x missing=0x%02x resident_submit=%.3f ms "
                                    "missing_bind=%.3f ms\n",
                                    layer_index,
                                    stream_expert_resident_mask,
                                    stream_expert_missing_mask,
                                    stream_split_resident_ms,
                                    stream_split_missing_ms);
                        }
                    } else {
                        ds4_gpu_stream_expert_split_args split_args = {
                            .active_mask = 0x3fu,
                            .accumulate = 0u,
                        };
                        ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb,
                                                                               g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline,
                                                                               &gate_args,
                                                                               &act_args,
                                                                               &split_args,
                                                                               stream_slot_entries,
                                                                               stream_gate_addr_buf,
                                                                               stream_up_addr_buf,
                                                                               xbuf,
                                                                               ds4_gpu_tensor_offset(x),
                                                                               gatebuf,
                                                                               ds4_gpu_tensor_offset(gate),
                                                                               upbuf,
                                                                               ds4_gpu_tensor_offset(up),
                                                                               midbuf,
                                                                               ds4_gpu_tensor_offset(mid),
                                                                               selected_exec_buf,
                                                                               selected_exec_off,
                                                                               weightsbuf,
                                                                               ds4_gpu_tensor_offset(weights),
                                                                               gate_smem,
                                                                               2,
                                                                               false);
                    }
                } else {
                    ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(cb,
                                                                    g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline,
                                                                    &gate_args,
                                                                    &act_args,
                                                                    stream_slot_entries,
                                                                    n_expert,
                                                                    stream_gate_addr_buf,
                                                                    stream_up_addr_buf,
                                                                    xbuf,
                                                                    ds4_gpu_tensor_offset(x),
                                                                    gatebuf,
                                                                    ds4_gpu_tensor_offset(gate),
                                                                    upbuf,
                                                                    ds4_gpu_tensor_offset(up),
                                                                    midbuf,
                                                                    ds4_gpu_tensor_offset(mid),
                                                                    selected_exec_buf,
                                                                    selected_exec_off,
                                                                    weightsbuf,
                                                                    ds4_gpu_tensor_offset(weights),
                                                                    gate_smem,
                                                                    2,
                                                                    false,
                                                                    nil,
                                                                    nil);
                }
            } else {
                ok = (!use_stream_expert_cache ||
                      ds4_gpu_stream_expert_cache_mark_entries_inflight(
                              stream_slot_entries,
                              n_expert,
                              0)) &&
                     ds4_gpu_encode_mul_mv_slots6_pair_swiglu(cb,
                                                              slots_pair_swiglu_pipeline,
                                                              &gate_args,
                                                              &act_args,
                                                              gate_slot_bufs,
                                                              gate_slot_offsets,
                                                              up_slot_bufs,
                                                              up_slot_offsets,
                                                              xbuf,
                                                              ds4_gpu_tensor_offset(x),
                                                              gatebuf,
                                                              ds4_gpu_tensor_offset(gate),
                                                              upbuf,
                                                              ds4_gpu_tensor_offset(up),
                                                              midbuf,
                                                              ds4_gpu_tensor_offset(mid),
                                                              weightsbuf,
                                                              ds4_gpu_tensor_offset(weights),
                                                              gate_smem,
                                                              2,
                                                              false);
            }
        } else if (fuse_pair_swiglu) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb,
                                                        pair_swiglu_pipeline,
                                                        &gate_args,
                                                        &act_args,
                                                        gate_buf,
                                                        (NSUInteger)gate_inner,
                                                        up_buf,
                                                        (NSUInteger)up_inner,
                                                        xbuf,
                                                        ds4_gpu_tensor_offset(x),
                                                        gatebuf,
                                                        ds4_gpu_tensor_offset(gate),
                                                        upbuf,
                                                        ds4_gpu_tensor_offset(up),
                                                        midbuf,
                                                        ds4_gpu_tensor_offset(mid),
                                                        selectedbuf,
                                                        ds4_gpu_tensor_offset(selected),
                                                        weightsbuf,
                                                        ds4_gpu_tensor_offset(weights),
                                                        gate_smem,
                                                        2,
                                                        false);
        } else if (!g_quality_mode &&
                   gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
                   g_moe_mul_mv_id_iq2_xxs_pair_pipeline) {
            ok = ds4_gpu_encode_mul_mv_id_pair(cb,
                                                 g_moe_mul_mv_id_iq2_xxs_pair_pipeline,
                                                 &gate_args,
                                                 gate_buf,
                                                 (NSUInteger)gate_inner,
                                                 up_buf,
                                                 (NSUInteger)up_inner,
                                                 xbuf,
                                                 ds4_gpu_tensor_offset(x),
                                                 gatebuf,
                                                 ds4_gpu_tensor_offset(gate),
                                                 upbuf,
                                                 ds4_gpu_tensor_offset(up),
                                                 selectedbuf,
                                                 ds4_gpu_tensor_offset(selected),
                                                 gate_smem,
                                                 2,
                                                 false);
        } else if (!g_quality_mode &&
                   gate_type == DS4_METAL_TENSOR_Q4_K &&
                   g_moe_mul_mv_id_q4_k_pair_pipeline) {
            ok = ds4_gpu_encode_mul_mv_id_pair(cb,
                                                 g_moe_mul_mv_id_q4_k_pair_pipeline,
                                                 &gate_args,
                                                 gate_buf,
                                                 (NSUInteger)gate_inner,
                                                 up_buf,
                                                 (NSUInteger)up_inner,
                                                 xbuf,
                                                 ds4_gpu_tensor_offset(x),
                                                 gatebuf,
                                                 ds4_gpu_tensor_offset(gate),
                                                 upbuf,
                                                 ds4_gpu_tensor_offset(up),
                                                 selectedbuf,
                                                 ds4_gpu_tensor_offset(selected),
                                                 gate_smem,
                                                 2,
                                                 false);
        } else {
            ok = ds4_gpu_encode_mul_mv_id(cb,
                                            gate_mv_pipeline,
                                            &gate_args,
                                            gate_buf,
                                            (NSUInteger)gate_inner,
                                            xbuf,
                                            ds4_gpu_tensor_offset(x),
                                            gatebuf,
                                            ds4_gpu_tensor_offset(gate),
                                            selectedbuf,
                                            ds4_gpu_tensor_offset(selected),
                                            gate_smem,
                                            gate_nsg,
                                            gate_rows_per_group_is_nr0) &&
                 ds4_gpu_encode_mul_mv_id(cb,
                                            gate_mv_pipeline,
                                            &gate_args,
                                            up_buf,
                                            (NSUInteger)up_inner,
                                            xbuf,
                                            ds4_gpu_tensor_offset(x),
                                            upbuf,
                                            ds4_gpu_tensor_offset(up),
                                            selectedbuf,
                                            ds4_gpu_tensor_offset(selected),
                                            gate_smem,
                                            gate_nsg,
                                            gate_rows_per_group_is_nr0);
        }
        DS4_METAL_PROFILE_MOE_ONE_STAGE("gate_up");
        if (ok && (!fuse_pair_swiglu || use_q4_group24_experts)) {
            ok = ds4_gpu_encode_moe_swiglu_weight(cb,
                                                    gatebuf,
                                                    ds4_gpu_tensor_offset(gate),
                                                    upbuf,
                                                    ds4_gpu_tensor_offset(up),
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    weightsbuf,
                                                    ds4_gpu_tensor_offset(weights),
                                                    expert_mid_dim,
                                                    pair_rows,
                                                    clamp,
                                                    false);
        }
        DS4_METAL_PROFILE_MOE_ONE_STAGE("activation_weight");

        id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer);
        NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) :
            (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0);
        if (ok && stream_expert_split_completed) {
            /* The split path already wrote the resident partial output and
             * accumulated the missing experts into out. */
        } else if (ok && use_q4_grouped_experts) {
            bool first_group = true;
            for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) {
                const uint32_t expert_count =
                    q4_expert_group_size < n_total_expert - expert_base ?
                    q4_expert_group_size : n_total_expert - expert_base;
                if ((uint64_t)expert_base > UINT64_MAX / down_expert_bytes ||
                    (uint64_t)expert_count > UINT64_MAX / down_expert_bytes) {
                    ok = 0;
                    break;
                }
                const uint64_t group_rel = (uint64_t)expert_base * down_expert_bytes;
                const uint64_t group_bytes = (uint64_t)expert_count * down_expert_bytes;
                if (group_rel > UINT64_MAX - down_offset) {
                    ok = 0;
                    break;
                }
                uint64_t down_group_inner = 0;
                id<MTLBuffer> down_group_buf =
                    q4_grouped_cache_views ?
                    ds4_gpu_wrap_model_exact_range(model_map,
                                                   model_size,
                                                   down_offset + group_rel,
                                                   group_bytes,
                                                   &down_group_inner) :
                    ds4_gpu_wrap_model_exact_range_transient(model_map,
                                                             model_size,
                                                             down_offset + group_rel,
                                                             group_bytes,
                                                             &down_group_inner);
                if (!down_group_buf) {
                    ok = 0;
                    break;
                }
                ds4_gpu_moe_expert_group_args group_args = {
                    .expert_base = expert_base,
                    .expert_count = expert_count,
                    .accumulate = first_group ? 0u : 1u,
                    .pad0 = 0,
                };
                ok = ds4_gpu_encode_mul_mv_group_q4_sum6(cb,
                                                          g_moe_mul_mv_group_q4_k_sum6_pipeline,
                                                          &down_args,
                                                          &group_args,
                                                          down_group_buf,
                                                          (NSUInteger)down_group_inner,
                                                          midbuf,
                                                          ds4_gpu_tensor_offset(mid),
                                                          outbuf,
                                                          ds4_gpu_tensor_offset(out),
                                                          selectedbuf,
                                                          ds4_gpu_tensor_offset(selected),
                                                          down_smem,
                                                          2);
                first_group = false;
            }
        } else if (ok && use_q4_expert_address_table) {
            ok = ds4_gpu_encode_mul_mv_addr_q4_sum6(cb,
                                                    g_moe_mul_mv_addr_q4_k_sum6_pipeline,
                                                    &down_args,
                                                    down_table,
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    outbuf,
                                                    ds4_gpu_tensor_offset(out),
                                                    selectedbuf,
                                                    ds4_gpu_tensor_offset(selected),
                                                    down_smem,
                                                    2);
        } else if (ok && use_q4_expert_table) {
            ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb,
                                                    g_moe_mul_mv_table_q4_k_sum6_pipeline,
                                                    &down_args,
                                                    down_table,
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    outbuf,
                                                    ds4_gpu_tensor_offset(out),
                                                    selectedbuf,
                                                    ds4_gpu_tensor_offset(selected),
                                                    down_smem,
                                                    2,
                                                    q4_table_queue_residency);
        } else if (ok && use_q4_group6_experts) {
            ok = ds4_gpu_encode_mul_mv_group6_sum6(cb,
                                                    g_moe_mul_mv_group6_q4_k_sum6_pipeline,
                                                    &down_args,
                                                    down_group6_bufs,
                                                    down_group6_offsets,
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    outbuf,
                                                    ds4_gpu_tensor_offset(out),
                                                    selectedbuf,
                                                    ds4_gpu_tensor_offset(selected),
                                                    down_smem,
                                                    2);
        } else if (ok && use_q4_group8_experts) {
            ok = ds4_gpu_encode_mul_mv_group8_sum6(cb,
                                                    g_moe_mul_mv_group8_q4_k_sum6_pipeline,
                                                    &down_args,
                                                    down_group8_bufs,
                                                    down_group8_offsets,
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    outbuf,
                                                    ds4_gpu_tensor_offset(out),
                                                    selectedbuf,
                                                    ds4_gpu_tensor_offset(selected),
                                                    down_smem,
                                                    2);
        } else if (ok && use_q4_group24_experts) {
            ok = ds4_gpu_encode_mul_mv_group24_sum6(cb,
                                                    g_moe_mul_mv_group24_q4_k_sum6_pipeline,
                                                    &down_args,
                                                    down_group24_bufs,
                                                    down_group24_offsets,
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    outbuf,
                                                    ds4_gpu_tensor_offset(out),
                                                    selectedbuf,
                                                    ds4_gpu_tensor_offset(selected),
                                                    down_smem,
                                                    2);
        } else if (ok && (use_q4_gather_slots || use_selected_slots)) {
            if (use_stream_expert_addr_table) {
                if (down_type == DS4_METAL_TENSOR_IQ2_XXS) {
                    ok = ds4_gpu_encode_mul_mv_addr_iq2(cb,
                                                        g_moe_mul_mv_addr_iq2_xxs_pipeline,
                                                        &down_args,
                                                        stream_slot_entries,
                                                        n_expert,
                                                        stream_down_addr_buf,
                                                        midbuf,
                                                        ds4_gpu_tensor_offset(mid),
                                                        down_dst,
                                                        down_dst_off,
                                                        selected_exec_buf,
                                                        selected_exec_off,
                                                        down_smem,
                                                        2,
                                                        false);
                } else if (use_stream_expert_masked_addr_table) {
                    ds4_gpu_stream_expert_split_args split_args = {
                        .active_mask = 0x3fu,
                        .accumulate = 0u,
                    };
                    ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb,
                                                                   g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline,
                                                                   &down_args,
                                                                   &split_args,
                                                                   stream_slot_entries,
                                                                   stream_down_addr_buf,
                                                                   midbuf,
                                                                   ds4_gpu_tensor_offset(mid),
                                                                   outbuf,
                                                                   ds4_gpu_tensor_offset(out),
                                                                   selected_exec_buf,
                                                                   selected_exec_off,
                                                                   down_smem,
                                                                   2);
                } else {
                    ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(cb,
                                                             g_moe_mul_mv_addr_q2_k_sum6_pipeline,
                                                             &down_args,
                                                             stream_slot_entries,
                                                             n_expert,
                                                             stream_down_addr_buf,
                                                             midbuf,
                                                             ds4_gpu_tensor_offset(mid),
                                                             outbuf,
                                                             ds4_gpu_tensor_offset(out),
                                                             selected_exec_buf,
                                                             selected_exec_off,
                                                             down_smem,
                                                             2,
                                                             nil);
                }
            } else {
                ok = (!use_stream_expert_cache ||
                      ds4_gpu_stream_expert_cache_mark_entries_inflight(
                              stream_slot_entries,
                              n_expert,
                              0)) &&
                     ds4_gpu_encode_mul_mv_slots6_sum6(cb,
                                                       slots_sum6_pipeline,
                                                       &down_args,
                                                       down_slot_bufs,
                                                       down_slot_offsets,
                                                       midbuf,
                                                       ds4_gpu_tensor_offset(mid),
                                                       outbuf,
                                                       ds4_gpu_tensor_offset(out),
                                                       down_smem,
                                                       2);
            }
        } else if (ok && direct_down_sum) {
            ok = ds4_gpu_encode_mul_mv_id_sum6(cb,
                                                 down_sum6_pipeline,
                                                 &down_args,
                                                 down_buf,
                                                 (NSUInteger)down_inner,
                                                 midbuf,
                                                 ds4_gpu_tensor_offset(mid),
                                                 outbuf,
                                                 ds4_gpu_tensor_offset(out),
                                                 selectedbuf,
                                                 ds4_gpu_tensor_offset(selected),
                                                 add_in ? ds4_gpu_tensor_buffer(add_in) : nil,
                                                 add_in ? ds4_gpu_tensor_offset(add_in) : 0,
                                                 down_smem,
                                                 2);
        } else if (ok) {
            ok = ds4_gpu_encode_mul_mv_id(cb,
                                                 down_mv_pipeline,
                                                 &down_args,
                                                 down_buf,
                                                 (NSUInteger)down_inner,
                                                 midbuf,
                                                 ds4_gpu_tensor_offset(mid),
                                                 down_dst,
                                                 down_dst_off,
                                                 selectedbuf,
                                                 ds4_gpu_tensor_offset(selected),
                                                 down_smem,
                                                 down_nsg,
                                                 down_rows_per_group_is_nr0);
        }
        DS4_METAL_PROFILE_MOE_ONE_STAGE("down");
        if (ok && n_expert > 1 && !direct_down_sum && !stream_expert_split_completed) {
            ok = ds4_gpu_encode_moe_sum_experts(cb,
                                                       down_dst,
                                                       down_dst_off,
                                                       outbuf,
                                                       ds4_gpu_tensor_offset(out),
                                                       out_dim,
                                                       n_expert,
                                                       n_tokens);
        }
        DS4_METAL_PROFILE_MOE_ONE_STAGE("sum");
        if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34395); return 0; }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34397); return 0; }
        if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34372);
                return 0;
            }
        }
#undef DS4_METAL_PROFILE_MOE_ONE_STAGE
    }

    return 1;
}

int ds4_gpu_routed_moe_batch_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        ds4_gpu_tensor       *experts,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                down_offset,
        uint32_t                gate_type,
        uint32_t                down_type,
        uint64_t                gate_expert_bytes,
        uint64_t                gate_row_bytes,
        uint64_t                down_expert_bytes,
        uint64_t                down_row_bytes,
        uint32_t                expert_in_dim,
        uint32_t                expert_mid_dim,
        uint32_t                out_dim,
        const ds4_gpu_tensor *selected,
        const ds4_gpu_tensor *weights,
        uint32_t                n_total_expert,
        uint32_t                n_expert,
        float                   clamp,
        const ds4_gpu_tensor *x,
        uint32_t                layer_index,
        uint32_t                n_tokens,
        bool                   *mid_is_f16,
        bool                    force_resident) {
    (void)force_resident;
    if (!g_initialized && !ds4_gpu_init()) return 0;
    /* TP sharding (see ds4_gpu_routed_moe_one_tensor): bind from the owned
     * expert range and rebase ids in the kernels. */
    uint32_t first_expert = 0;
    uint32_t n_bind_expert = 0;
    ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
    const int32_t tp_expert_base_host = (int32_t)first_expert;
    gate_offset += (uint64_t)first_expert * gate_expert_bytes;
    up_offset += (uint64_t)first_expert * gate_expert_bytes;
    down_offset += (uint64_t)first_expert * down_expert_bytes;
    if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights ||
        n_tokens == 0 || n_total_expert == 0 || n_expert == 0 ||
        n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED) {
        return 0;
    }
    if (gate_expert_bytes == 0 || down_expert_bytes == 0 ||
        gate_row_bytes == 0 || down_row_bytes == 0) {
        return 0;
    }
    if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) return 0;
    if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
        (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
        fprintf(stderr, "ds4: Metal routed batch MoE tensor byte size overflow\n");
        return 0;
    }
    const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;

    /*
     * PRO Q4 routed expert tensors are multi-GiB per layer.  A one-token
     * layer-slice prefill should use the one-token path so it can either bind
     * exact tensor views with GPU-selected IDs or fall back to selected-expert
     * views.  Keep this guarded by tensor size so Flash and mixed-Flash Q4 keep
     * their existing fast path.
     */
    const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull;
    const bool can_single_token_q4_grouped =
        getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL &&
        getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL &&
        g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil &&
        g_moe_mul_mv_group_q4_k_sum6_pipeline != nil;
    const bool can_single_token_q4_selected_slots =
        ds4_gpu_q4_selected_paths_allowed() &&
        getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL &&
        g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil &&
        g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil;
    const bool can_single_token_q4_group6 =
        n_total_expert == 384 &&
        getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL &&
        getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL &&
        g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil &&
        g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil;
    const bool can_single_token_q4_group8 =
        n_total_expert == 384 &&
        getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL &&
        getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL &&
        g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil &&
        g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil;
    const bool can_single_token_q4_group24 =
        n_total_expert == 384 &&
        getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL &&
        getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL &&
        g_moe_mul_mv_group24_q4_k_id_pipeline != nil &&
        g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil;
    const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0;
    const bool can_single_token_q4_exact_tensor_id =
        n_total_expert == 384 &&
        max_buffer_len != 0 &&
        gate_tensor_bytes <= max_buffer_len &&
        down_tensor_bytes <= max_buffer_len &&
        getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL &&
        getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL;
    const bool enable_single_token_q4_expert_table =
        getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
        ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
                                                 n_expert,
                                                 gate_tensor_bytes,
                                                 down_tensor_bytes);
    const bool can_single_token_q4_expert_table =
        n_total_expert == 384 &&
        enable_single_token_q4_expert_table &&
        getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL &&
        g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil &&
        g_moe_mul_mv_table_q4_k_sum6_pipeline != nil &&
        g_moe_table_q4_pair_gate_encoder != nil &&
        g_moe_table_q4_pair_up_encoder != nil &&
        g_moe_table_q4_sum_down_encoder != nil;
    const bool can_single_token_q4_expert_address_table =
        n_total_expert == 384 &&
        getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL &&
        getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL &&
        g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil &&
        g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil;
    const bool use_single_token_q4_one_tensor =
        gate_type == DS4_METAL_TENSOR_Q4_K &&
        down_type == DS4_METAL_TENSOR_Q4_K &&
        n_tokens == 1 &&
        n_expert == 6 &&
        n_total_expert >= 128 &&
        (g_ssd_streaming_mode ||
         (gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
          down_tensor_bytes >= q4_selected_min_tensor_bytes)) &&
        !g_quality_mode &&
        getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
        getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
        (can_single_token_q4_grouped ||
         can_single_token_q4_group6 ||
         can_single_token_q4_group8 ||
         can_single_token_q4_group24 ||
         can_single_token_q4_exact_tensor_id ||
         can_single_token_q4_expert_address_table ||
         can_single_token_q4_expert_table ||
         can_single_token_q4_selected_slots);
    if (use_single_token_q4_one_tensor) {
        if (mid_is_f16) *mid_is_f16 = false;
        return ds4_gpu_routed_moe_one_tensor(out,
                                             gate,
                                             up,
                                             mid,
                                             experts,
                                             model_map,
                                             model_size,
                                             gate_offset,
                                             up_offset,
                                             down_offset,
                                             gate_type,
                                             down_type,
                                             gate_expert_bytes,
                                             gate_row_bytes,
                                             down_expert_bytes,
                                             down_row_bytes,
                                             expert_in_dim,
                                             expert_mid_dim,
                                             out_dim,
                                             selected,
                                             weights,
                                             n_total_expert,
                                             n_expert,
                                             clamp,
                                             x,
                                             NULL,
                                             layer_index,
                                             false);
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
        id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> expertsbuf = ds4_gpu_tensor_buffer(experts);
        id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
        id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
        const uint64_t x_bytes = (uint64_t)n_tokens * expert_in_dim * sizeof(float);
        const uint64_t mid_bytes = (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float);
        const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float);
        const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert * sizeof(int);
        const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert * sizeof(float);
        if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(gate) < mid_bytes ||
            ds4_gpu_tensor_bytes(up) < mid_bytes ||
            ds4_gpu_tensor_bytes(mid) < mid_bytes ||
            ds4_gpu_tensor_bytes(out) < out_bytes ||
            ds4_gpu_tensor_bytes(selected) < selected_bytes ||
            ds4_gpu_tensor_bytes(weights) < weights_bytes) {
            fprintf(stderr, "ds4: Metal routed batch MoE received undersized activation buffers\n");
            return 0;
        }
        uint64_t gate_inner = 0;
        uint64_t up_inner = 0;
        uint64_t down_inner = 0;
        id<MTLBuffer> gate_buf = nil;
        id<MTLBuffer> up_buf = nil;
        id<MTLBuffer> down_buf = nil;
        DS4MetalQ4ExpertTable *gate_table = nil;
        DS4MetalQ4ExpertTable *up_table = nil;
        DS4MetalQ4ExpertTable *down_table = nil;
        id q4_table_layer_residency = nil;
        id<MTLBuffer> stream_gate_addr_buf = nil;
        id<MTLBuffer> stream_up_addr_buf = nil;
        id<MTLBuffer> stream_down_addr_buf = nil;
        id<MTLBuffer> stream_overflow_gate = nil;
        id<MTLBuffer> stream_overflow_up = nil;
        id<MTLBuffer> stream_overflow_down = nil;
        ds4_gpu_stream_expert_cache_entry
            *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL };
        uint32_t stream_resource_count = 0;
        uint32_t stream_unique = 0;

        const uint32_t pair_rows = n_tokens * n_expert;
        const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float);

        const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type);
        const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type);
        id<MTLComputePipelineState> gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type);
        id<MTLComputePipelineState> down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type);
        id<MTLComputePipelineState> gate_mm_pipeline = nil;
        id<MTLComputePipelineState> up_mm_pipeline = nil;
        id<MTLComputePipelineState> down_mm_pipeline = nil;
        id<MTLComputePipelineState> pair_swiglu_mm_pipeline = nil;
        if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) {
            fprintf(stderr, "ds4: unsupported Metal routed batch MoE quant types gate=%u down=%u\n",
                    gate_type, down_type);
            return 0;
        }
        const bool use_iq2_batch_selected_addr =
            ds4_gpu_stream_prefill_batch_selected_addr_enabled(n_tokens,
                                                               n_total_expert,
                                                               n_expert,
                                                               gate_type,
                                                               down_type) &&
            n_tokens > 1 &&
            n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
            !g_quality_mode &&
            getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
            getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
            g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil;

        ds4_gpu_mul_mv_id_args gate_args =
            ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                          gate_row_bytes, gate_expert_bytes,
                                          1, n_expert, n_tokens, gate_nr0);
        gate_args.tp_rank = g_tp_split_rank;
        gate_args.tp_world = g_tp_split_world;
        gate_args.tp_expert_base = tp_expert_base_host;
        ds4_gpu_mul_mv_id_args down_args =
            ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert,
                                          down_row_bytes, down_expert_bytes,
                                          n_expert, n_expert, n_tokens, down_nr0);
        down_args.tp_rank = g_tp_split_rank;
        down_args.tp_world = g_tp_split_world;
        down_args.tp_expert_base = tp_expert_base_host;
        const bool q4_batch_expert_table_auto =
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
                                                     n_expert,
                                                     gate_tensor_bytes,
                                                     down_tensor_bytes);
        const bool q4_batch_table_queue_residency =
            ds4_gpu_q4_table_queue_residency_enabled(q4_batch_expert_table_auto);
        const bool enable_q4_batch_expert_table =
            getenv("DS4_METAL_ENABLE_Q4_BATCH_EXPERT_TABLE") != NULL ||
            getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
            q4_batch_expert_table_auto;
        const bool use_q4_batch_expert_table =
            gate_type == DS4_METAL_TENSOR_Q4_K &&
            down_type == DS4_METAL_TENSOR_Q4_K &&
            n_expert == 6 &&
            n_tokens > 1 &&
            n_total_expert == 384 &&
            gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
            down_tensor_bytes >= q4_selected_min_tensor_bytes &&
            !g_quality_mode &&
            getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
            getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
            getenv("DS4_METAL_DISABLE_Q4_BATCH_EXPERT_TABLE") == NULL &&
            getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL &&
            g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil &&
            g_moe_mul_mv_table_q4_k_sum6_pipeline != nil &&
            g_moe_table_q4_pair_gate_encoder != nil &&
            g_moe_table_q4_pair_up_encoder != nil &&
            g_moe_table_q4_sum_down_encoder != nil &&
            enable_q4_batch_expert_table &&
            (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL ||
             getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL ||
             q4_batch_table_queue_residency ||
             ds4_gpu_q4_table_model_residency_enabled());
        const bool use_mm_id =
            !use_q4_batch_expert_table &&
            !use_iq2_batch_selected_addr &&
            n_tokens >= 32u &&
            ds4_gpu_mul_mm_id_map0_name(n_expert) != NULL;
        /*
         * MTP verification is neither normal decode nor large prefill: the
         * target model must verify a tiny suffix (up to DSpark's 5-token
         * block) in one layer-major pass.  For that shape the prefill
         * expert-major GEMM path
         * is too large, but the decode pair kernels are exactly the right
         * primitive: they read the same activation once and compute routed
         * gate/up together for every selected expert row.  Keep this limited to
         * tiny batches so ordinary prefill keeps using the higher-throughput
         * grouped matmul path.
         */
        const bool use_tiny_pair_mv =
            !g_quality_mode &&
            n_tokens <= 5u &&
            !use_q4_batch_expert_table &&
            !use_mm_id &&
            ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) ||
             (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline));
        id<MTLComputePipelineState> tiny_pair_swiglu_pipeline = nil;
        if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) {
            tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline;
        } else if (gate_type == DS4_METAL_TENSOR_Q4_K) {
            tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline;
        }
        const bool use_tiny_pair_swiglu =
            use_tiny_pair_mv &&
            tiny_pair_swiglu_pipeline != nil &&
            getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL &&
            getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL;
        ds4_gpu_mul_mm_id_map_args gate_map_args = { 0 };
        ds4_gpu_mul_mm_id_args gate_mm_args = { 0 };
        ds4_gpu_mul_mm_id_args down_mm_args = { 0 };
        id<MTLComputePipelineState> map_pipeline = nil;
        /*
         * The grouped routed-MoE matmul loads activation tiles as half before
         * using SIMD-group MMA.  Store the SwiGLU/route-weight intermediate in
         * that same precision so the down projection avoids a large F32 mid
         * write/read. --quality keeps the older F32 intermediate.
         */
        const bool request_mid_f16 =
            !g_quality_mode &&
            !use_q4_batch_expert_table &&
            !use_iq2_batch_selected_addr;
        /*
         * Fused gate+up grouped matmul with the SwiGLU epilogue.  The IQ2
         * variant stays opt-in behind its env flag; the Q4_K variant is the
         * default path — same MMA accumulation order and epilogue math as the
         * separate GEMMs + swiglu pass, so the mid tensor is bit-identical.
         */
        const bool use_mm_id_pair_swiglu =
            use_mm_id &&
            g_tp_split_world != 2 &&    /* pair-swiglu mm kernel lacks expert ownership */
            request_mid_f16 &&
            n_expert == 6 &&
            ((gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
              down_type == DS4_METAL_TENSOR_Q2_K &&
              getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL) ||
             (gate_type == DS4_METAL_TENSOR_Q4_K &&
              down_type == DS4_METAL_TENSOR_Q4_K)) &&
            getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL &&
            getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
            getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL;
        if (use_mm_id) {
            gate_map_args =
                ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens);
            gate_mm_args =
                ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
                                              gate_row_bytes, gate_expert_bytes,
                                              1, n_expert, n_tokens);
            down_mm_args =
                ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert,
                                                        down_row_bytes, down_expert_bytes,
                                                        n_expert, n_expert, n_tokens,
                                                        request_mid_f16 ? sizeof(uint16_t) : sizeof(float));
            gate_mm_args.tp_rank = g_tp_split_rank;
            gate_mm_args.tp_world = g_tp_split_world;
            gate_mm_args.tp_expert_base = tp_expert_base_host;
            down_mm_args.tp_rank = g_tp_split_rank;
            down_mm_args.tp_world = g_tp_split_world;
            down_mm_args.tp_expert_base = tp_expert_base_host;

            map_pipeline = ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert));
            gate_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type);
            up_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type);
            down_mm_pipeline = request_mid_f16 ?
                ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) :
                ds4_gpu_routed_mm_pipeline(down_type);
            if (use_mm_id_pair_swiglu) {
                pair_swiglu_mm_pipeline =
                    ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_Q4_K ?
                        "kernel_mul_mm_id_q4_K_pair_swiglu_f16" :
                        "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16");
            }
            if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline ||
                (use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) {
                return 0;
            }
        }

        if (use_iq2_batch_selected_addr) {
            const int had_batch = g_batch_cb != nil;
            if (had_batch && ds4_gpu_end_commands() == 0) {
                return 0;
            }
            g_stream_prefill_batch_selected_addr_building++;
            if (!ds4_gpu_stream_expert_cache_prepare_selected_batch(
                        model_map,
                        model_size,
                        layer_index,
                        selected,
                        n_tokens,
                        n_total_expert,
                        n_expert,
                        gate_offset,
                        up_offset,
                        down_offset,
                        gate_expert_bytes,
                        down_expert_bytes,
                        &stream_gate_addr_buf,
                        &stream_up_addr_buf,
                        &stream_down_addr_buf,
                        stream_resources,
                        &stream_resource_count,
                        &stream_unique,
                        &stream_overflow_gate,
                        &stream_overflow_up,
                        &stream_overflow_down)) {
                g_stream_prefill_batch_selected_addr_building--;
                return 0;
            }
            g_stream_prefill_batch_selected_addr_building--;
            if (stream_unique == 0) {
                ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                fprintf(stderr,
                        "ds4: Metal streaming prefill batch selected addr layer=%u "
                        "produced no resident experts\n",
                        layer_index);
                return 0;
            }
            for (uint32_t i = 0; i < stream_resource_count; i++) {
                ds4_gpu_stream_expert_cache_entry *entry = stream_resources[i];
                if (!entry ||
                    !entry->valid ||
                    !entry->gate_buffer ||
                    !entry->up_buffer ||
                    !entry->down_buffer) {
                    fprintf(stderr,
                            "ds4: Metal streaming prefill batch selected addr layer=%u "
                            "lost resident expert resource %u/%u during preparation\n",
                            layer_index,
                            i,
                            stream_resource_count);
                    ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                    return 0;
                }
            }
            if (had_batch && ds4_gpu_begin_commands() == 0) {
                fprintf(stderr,
                        "ds4: Metal streaming prefill batch selected addr layer=%u "
                        "failed to reopen command batch after preparation\n",
                        layer_index);
                ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                return 0;
            }
        }

        if (use_q4_batch_expert_table) {
            gate_table = ds4_gpu_q4_expert_table(model_map,
                                                 model_size,
                                                 gate_offset,
                                                 gate_expert_bytes,
                                                 n_total_expert,
                                                 g_moe_table_q4_pair_gate_encoder);
            up_table = ds4_gpu_q4_expert_table(model_map,
                                               model_size,
                                               up_offset,
                                               gate_expert_bytes,
                                               n_total_expert,
                                               g_moe_table_q4_pair_up_encoder);
            down_table = ds4_gpu_q4_expert_table(model_map,
                                                 model_size,
                                                 down_offset,
                                                 down_expert_bytes,
                                                 n_total_expert,
                                                 g_moe_table_q4_sum_down_encoder);
            if (!gate_table || !up_table || !down_table) {
                return 0;
            }
            q4_table_layer_residency =
                ds4_gpu_q4_expert_layer_residency_set(gate_table,
                                                       up_table,
                                                       down_table,
                                                       q4_batch_expert_table_auto);
            if ((q4_batch_table_queue_residency ||
                 getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL) &&
                !q4_table_layer_residency &&
                !ds4_gpu_q4_table_model_residency_enabled()) {
                fprintf(stderr, "ds4: Metal Q4 batch expert table residency set is not available\n");
                return 0;
            }
        } else if (use_iq2_batch_selected_addr) {
            if (n_expert > 1 && (!expertsbuf ||
                ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) {
                if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
                                                   &g_moe_down_scratch_bytes,
                                                   (NSUInteger)down_scratch_bytes,
                                                   "ds4_moe_down_scratch")) {
                    ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                    return 0;
                }
            }
        } else {
            if (n_expert > 1 && (!expertsbuf ||
                ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) {
                if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
                                                   &g_moe_down_scratch_bytes,
                                                   (NSUInteger)down_scratch_bytes,
                                                   "ds4_moe_down_scratch")) {
                    return 0;
                }
            }
            gate_buf = ds4_gpu_wrap_model_range(model_map,
                                                model_size,
                                                gate_offset,
                                                gate_tensor_bytes,
                                                &gate_inner);
            up_buf = ds4_gpu_wrap_model_range(model_map,
                                              model_size,
                                              up_offset,
                                              gate_tensor_bytes,
                                              &up_inner);
            down_buf = ds4_gpu_wrap_model_range(model_map,
                                                model_size,
                                                down_offset,
                                                down_tensor_bytes,
                                                &down_inner);
            if (!gate_buf || !up_buf || !down_buf) return 0;
            if (getenv("DS4_GLM_TP_DEBUG") && layer_index == 3) {
                fprintf(stderr,
                        "ds4: batch mv binds l=%u rank=%d base=%d gate=%llu+%llu "
                        "up=%llu down=%llu inner=%llu/%llu/%llu ne02=%d nei0=%d nr0=%d\n",
                        layer_index, g_tp_split_rank, tp_expert_base_host,
                        (unsigned long long)gate_offset,
                        (unsigned long long)gate_tensor_bytes,
                        (unsigned long long)up_offset,
                        (unsigned long long)down_offset,
                        (unsigned long long)gate_inner,
                        (unsigned long long)up_inner,
                        (unsigned long long)down_inner,
                        gate_args.ne02, gate_args.nei0, gate_args.nr0);
            }
        }

        const bool q4_batch_table_boundary =
            use_q4_batch_expert_table &&
            g_batch_cb != nil &&
            getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL &&
            !q4_batch_table_queue_residency &&
            getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL;
        if (q4_batch_table_boundary) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;
        if (use_q4_batch_expert_table && !ds4_gpu_use_model_residency_set(cb)) {
            return 0;
        }
        if (!q4_batch_table_queue_residency &&
            q4_table_layer_residency &&
            [cb respondsToSelector:@selector(useResidencySet:)]) {
            [cb useResidencySet:q4_table_layer_residency];
        }
        const bool moe_stage_profile =
            g_batch_cb != nil &&
            ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_STAGE_PROFILE",
                                                    "DS4_METAL_MOE_STAGE_PROFILE_LAYER",
                                                    layer_index);
        const char *moe_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER");
        const char *moe_path =
            use_q4_batch_expert_table ? "q4_table_pair_swiglu" :
            use_iq2_batch_selected_addr ? "iq2_batch_stream_addr" :
            use_mm_id_pair_swiglu ? "mm_id_pair_swiglu" :
            use_mm_id ? "mm_id" :
            use_tiny_pair_swiglu ? "tiny_pair_swiglu" :
            (use_tiny_pair_mv ? "tiny_pair_mv" : "mv");
        double moe_stage_t0 = moe_stage_profile ? ds4_gpu_now_ms() : 0.0;
        if (moe_stage_profile) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;
            moe_stage_t0 = ds4_gpu_now_ms();
        }
#define DS4_METAL_PROFILE_MOE_STAGE(name) do { \
            if (ok && moe_stage_profile) { \
                if (ds4_gpu_end_commands() == 0) { \
                    ok = 0; \
                } else { \
                    const char *stage_name = (name); \
                    const double now_ms = ds4_gpu_now_ms(); \
                    const int print_stage = \
                        !moe_stage_filter || !moe_stage_filter[0] || \
                        strstr(stage_name, moe_stage_filter) != NULL; \
                    if (print_stage) { \
                        fprintf(stderr, \
                                "ds4: Metal routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \
                                "gate=%s down=%s path=%s mid=%s %s=%.3f ms\n", \
                                layer_index, n_tokens, pair_rows, n_expert, \
                                ds4_gpu_metal_tensor_type_name(gate_type), \
                                ds4_gpu_metal_tensor_type_name(down_type), \
                                moe_path, \
                                request_mid_f16 ? "f16" : "f32", \
                                stage_name, now_ms - moe_stage_t0); \
                    } \
                    moe_stage_t0 = now_ms; \
                    if (ds4_gpu_begin_commands() == 0) { \
                        ok = 0; \
                    } else { \
                        cb = ds4_gpu_command_buffer(&owned); \
                        if (!cb) ok = 0; \
                    } \
                } \
            } \
        } while (0)

        const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type);
        const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type);
        const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type);
        const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type);
        const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type);
        const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type);
        id<MTLComputePipelineState> down_sum6_pipeline = nil;
        if (down_type == DS4_METAL_TENSOR_Q2_K) {
            down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline;
        } else if (down_type == DS4_METAL_TENSOR_Q4_K) {
            down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline;
        }
        const bool direct_down_sum =
            !g_quality_mode &&
            !use_q4_batch_expert_table &&
            !use_mm_id &&
            n_expert == 6 &&
            n_tokens <= 4u &&
            down_sum6_pipeline != nil;
        int ok = 0;
        if (use_iq2_batch_selected_addr) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(
                    cb,
                    g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline,
                    &gate_args,
                    &act_args,
                    stream_resources,
                    stream_resource_count,
                    stream_gate_addr_buf,
                    stream_up_addr_buf,
                    xbuf,
                    ds4_gpu_tensor_offset(x),
                    gatebuf,
                    ds4_gpu_tensor_offset(gate),
                    upbuf,
                    ds4_gpu_tensor_offset(up),
                    midbuf,
                    ds4_gpu_tensor_offset(mid),
                    selectedbuf,
                    ds4_gpu_tensor_offset(selected),
                    weightsbuf,
                    ds4_gpu_tensor_offset(weights),
                    gate_smem,
                    2,
                    false,
                    stream_overflow_gate,
                    stream_overflow_up);
            if (!ok) {
                fprintf(stderr,
                        "ds4: Metal streaming prefill batch selected addr layer=%u "
                        "failed to encode gate/up path tokens=%u unique=%u\n",
                        layer_index,
                        n_tokens,
                        stream_unique);
            }
        } else if (use_q4_batch_expert_table) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb,
                                                            g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline,
                                                            &gate_args,
                                                            &act_args,
                                                            gate_table,
                                                            up_table,
                                                            xbuf,
                                                            ds4_gpu_tensor_offset(x),
                                                            gatebuf,
                                                            ds4_gpu_tensor_offset(gate),
                                                            upbuf,
                                                            ds4_gpu_tensor_offset(up),
                                                            midbuf,
                                                            ds4_gpu_tensor_offset(mid),
                                                            selectedbuf,
                                                            ds4_gpu_tensor_offset(selected),
                                                            weightsbuf,
                                                            ds4_gpu_tensor_offset(weights),
                                                            gate_smem,
                                                            2,
                                                            false,
                                                            q4_batch_table_queue_residency);
        } else if (use_mm_id) {
            /*
             * The routed pair ids are the same for gate, up, and down. Build
             * the expert-major work map once, then reuse it for all three
             * batched expert matmuls.
             */
            ok = ds4_gpu_encode_mul_mm_id_map(cb,
                                                map_pipeline,
                                                &gate_map_args,
                                                &gate_mm_args,
                                                selectedbuf,
                                                ds4_gpu_tensor_offset(selected));
            DS4_METAL_PROFILE_MOE_STAGE("map");
            if (ok && use_mm_id_pair_swiglu) {
                ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                    .width = expert_mid_dim,
                    .rows = pair_rows,
                    .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                    .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                    .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(uint16_t),
                    .weight_stride = sizeof(float),
                    .write_clamped = 0,
                    .clamp_value = clamp,
                };
                ok = ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(cb,
                                                                   pair_swiglu_mm_pipeline,
                                                                   &gate_mm_args,
                                                                   &act_args,
                                                                   gate_buf,
                                                                   (NSUInteger)gate_inner,
                                                                   up_buf,
                                                                   (NSUInteger)up_inner,
                                                                   xbuf,
                                                                   ds4_gpu_tensor_offset(x),
                                                                   midbuf,
                                                                   ds4_gpu_tensor_offset(mid),
                                                                   weightsbuf,
                                                                   ds4_gpu_tensor_offset(weights));
                DS4_METAL_PROFILE_MOE_STAGE("gate_up_fused");
            } else if (ok) {
                ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                           gate_mm_pipeline,
                                                           &gate_mm_args,
                                                           gate_buf,
                                                           (NSUInteger)gate_inner,
                                                           xbuf,
                                                           ds4_gpu_tensor_offset(x),
                                                           gatebuf,
                                                           ds4_gpu_tensor_offset(gate),
                                                           8192u);
                DS4_METAL_PROFILE_MOE_STAGE("gate");
            }
            if (ok && !use_mm_id_pair_swiglu) {
                ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                   up_mm_pipeline,
                                                   &gate_mm_args,
                                                   up_buf,
                                                   (NSUInteger)up_inner,
                                                   xbuf,
                                                   ds4_gpu_tensor_offset(x),
                                                   upbuf,
                                                   ds4_gpu_tensor_offset(up),
                                                   8192u);
                DS4_METAL_PROFILE_MOE_STAGE("up");
            }
        } else if (use_tiny_pair_swiglu) {
            ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
                .width = expert_mid_dim,
                .rows = pair_rows,
                .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
                .weight_stride = sizeof(float),
                .write_clamped = 0,
                .clamp_value = clamp,
            };
            ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb,
                                                       tiny_pair_swiglu_pipeline,
                                                       &gate_args,
                                                       &act_args,
                                                       gate_buf,
                                                       (NSUInteger)gate_inner,
                                                       up_buf,
                                                       (NSUInteger)up_inner,
                                                       xbuf,
                                                       ds4_gpu_tensor_offset(x),
                                                       gatebuf,
                                                       ds4_gpu_tensor_offset(gate),
                                                       upbuf,
                                                       ds4_gpu_tensor_offset(up),
                                                       midbuf,
                                                       ds4_gpu_tensor_offset(mid),
                                                       selectedbuf,
                                                       ds4_gpu_tensor_offset(selected),
                                                       weightsbuf,
                                                       ds4_gpu_tensor_offset(weights),
                                                       gate_smem,
                                                       2,
                                                       false);
        } else if (use_tiny_pair_mv) {
            id<MTLComputePipelineState> pair_pipeline =
                gate_type == DS4_METAL_TENSOR_IQ2_XXS ?
                    g_moe_mul_mv_id_iq2_xxs_pair_pipeline :
                    g_moe_mul_mv_id_q4_k_pair_pipeline;
            ok = ds4_gpu_encode_mul_mv_id_pair(cb,
                                                 pair_pipeline,
                                                 &gate_args,
                                                 gate_buf,
                                                 (NSUInteger)gate_inner,
                                                 up_buf,
                                                 (NSUInteger)up_inner,
                                                 xbuf,
                                                 ds4_gpu_tensor_offset(x),
                                                 gatebuf,
                                                 ds4_gpu_tensor_offset(gate),
                                                 upbuf,
                                                 ds4_gpu_tensor_offset(up),
                                                 selectedbuf,
                                                 ds4_gpu_tensor_offset(selected),
                                                 gate_smem,
                                                 2,
                                                 false);
        } else {
            ok = ds4_gpu_encode_mul_mv_id(cb,
                                                  gate_mv_pipeline,
                                                  &gate_args,
                                                  gate_buf,
                                                  (NSUInteger)gate_inner,
                                                  xbuf,
                                                  ds4_gpu_tensor_offset(x),
                                                  gatebuf,
                                                  ds4_gpu_tensor_offset(gate),
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  gate_smem,
                                                  gate_nsg,
                                                  gate_rows_per_group_is_nr0) &&
                 ds4_gpu_encode_mul_mv_id(cb,
                                                  gate_mv_pipeline,
                                                  &gate_args,
                                                  up_buf,
                                                  (NSUInteger)up_inner,
                                                  xbuf,
                                                  ds4_gpu_tensor_offset(x),
                                                  upbuf,
                                                  ds4_gpu_tensor_offset(up),
                                                  selectedbuf,
                                                  ds4_gpu_tensor_offset(selected),
                                                  gate_smem,
                                                  gate_nsg,
                                                  gate_rows_per_group_is_nr0);
        }
        DS4_METAL_PROFILE_MOE_STAGE("gate_up");
        const bool use_fused_activation = !g_quality_mode && !use_q4_batch_expert_table;
        const bool use_mid_f16 =
            use_mm_id &&
            use_fused_activation &&
            request_mid_f16;
        if (mid_is_f16) *mid_is_f16 = use_mid_f16;
        if (ok && use_iq2_batch_selected_addr) {
            /* The address-table pair kernel already wrote weighted SwiGLU rows into mid. */
        } else if (ok && use_q4_batch_expert_table) {
            /* The table pair kernel already wrote weighted SwiGLU rows into mid. */
        } else if (ok && use_mm_id_pair_swiglu) {
            /* The fused batch mm_id pair kernel already wrote weighted f16 SwiGLU rows into mid. */
        } else if (ok && use_tiny_pair_swiglu) {
            /* The fused tiny pair kernel already wrote weighted F32 SwiGLU rows into mid. */
        } else if (ok && use_fused_activation) {
            ok = ds4_gpu_encode_moe_swiglu_weight(cb,
                                                    gatebuf,
                                                    ds4_gpu_tensor_offset(gate),
                                                    upbuf,
                                                    ds4_gpu_tensor_offset(up),
                                                    midbuf,
                                                    ds4_gpu_tensor_offset(mid),
                                                    weightsbuf,
                                                    ds4_gpu_tensor_offset(weights),
                                                    expert_mid_dim,
                                                    pair_rows,
                                                    clamp,
                                                    use_mid_f16);
        } else if (ok && clamp > 1.0e-6f) {
            ok = ds4_gpu_encode_unary_f32_rows(cb,
                                                 g_unary_clamp_pipeline,
                                                 gatebuf,
                                                 ds4_gpu_tensor_offset(gate),
                                                 gatebuf,
                                                 ds4_gpu_tensor_offset(gate),
                                                 expert_mid_dim,
                                                 pair_rows,
                                                 0,
                                                 -FLT_MAX,
                                                 clamp);
            if (ok) {
                ok = ds4_gpu_encode_unary_f32_rows(cb,
                                                     g_unary_silu_pipeline,
                                                     gatebuf,
                                                     ds4_gpu_tensor_offset(gate),
                                                     midbuf,
                                                     ds4_gpu_tensor_offset(mid),
                                                     expert_mid_dim,
                                                     pair_rows,
                                                     1,
                                                     0.0f,
                                                     0.0f);
            }
            if (ok) {
                ok = ds4_gpu_encode_unary_f32_rows(cb,
                                                 g_unary_clamp_pipeline,
                                                 upbuf,
                                                 ds4_gpu_tensor_offset(up),
                                                 upbuf,
                                                 ds4_gpu_tensor_offset(up),
                                                 expert_mid_dim,
                                                 pair_rows,
                                                 0,
                                                 -clamp,
                                                 clamp);
            }
            if (ok) {
                ds4_gpu_bin_args mul_args =
                    ds4_gpu_make_bin_same_rows_args(expert_mid_dim, pair_rows);
                ok = ds4_gpu_encode_bin_f32_rows(cb,
                                                   g_mul_pipeline,
                                                   &mul_args,
                                                   midbuf,
                                                   ds4_gpu_tensor_offset(mid),
                                                   upbuf,
                                                   ds4_gpu_tensor_offset(up),
                                                   midbuf,
                                                   ds4_gpu_tensor_offset(mid));
            }
        } else if (ok) {
            ok = ds4_gpu_encode_swiglu_flat(cb,
                                              gatebuf,
                                              ds4_gpu_tensor_offset(gate),
                                              upbuf,
                                              ds4_gpu_tensor_offset(up),
                                              midbuf,
                                              ds4_gpu_tensor_offset(mid),
                                              (uint32_t)((uint64_t)pair_rows * expert_mid_dim));
        }
        if (ok && !use_fused_activation && !use_q4_batch_expert_table) {
            ds4_gpu_bin_args weight_args =
                ds4_gpu_make_bin_rowwise_scalar_args(expert_mid_dim, pair_rows);
            ok = ds4_gpu_encode_bin_f32_rows(cb,
                                               g_bin_mul_scalar_pipeline,
                                               &weight_args,
                                               midbuf,
                                               ds4_gpu_tensor_offset(mid),
                                               weightsbuf,
                                               ds4_gpu_tensor_offset(weights),
                                               midbuf,
                                               ds4_gpu_tensor_offset(mid));
        }
        DS4_METAL_PROFILE_MOE_STAGE("activation_weight");

        id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer);
        NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) :
            (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0);
        if (ok) {
            if (use_iq2_batch_selected_addr) {
                ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(
                        cb,
                        g_moe_mul_mv_addr_q2_k_sum6_pipeline,
                        &down_args,
                        stream_resources,
                        stream_resource_count,
                        stream_down_addr_buf,
                        midbuf,
                        ds4_gpu_tensor_offset(mid),
                        outbuf,
                        ds4_gpu_tensor_offset(out),
                        selectedbuf,
                        ds4_gpu_tensor_offset(selected),
                        down_smem,
                        2,
                        stream_overflow_down);
                if (!ok) {
                    fprintf(stderr,
                            "ds4: Metal streaming prefill batch selected addr layer=%u "
                            "failed to encode down path tokens=%u unique=%u\n",
                            layer_index,
                            n_tokens,
                            stream_unique);
                }
            } else if (use_q4_batch_expert_table) {
                ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb,
                                                         g_moe_mul_mv_table_q4_k_sum6_pipeline,
                                                         &down_args,
                                                         down_table,
                                                         midbuf,
                                                         ds4_gpu_tensor_offset(mid),
                                                         outbuf,
                                                         ds4_gpu_tensor_offset(out),
                                                         selectedbuf,
                                                         ds4_gpu_tensor_offset(selected),
                                                         down_smem,
                                                         2,
                                                         q4_batch_table_queue_residency);
            } else if (direct_down_sum) {
                ok = ds4_gpu_encode_mul_mv_id_sum6(cb,
                                                     down_sum6_pipeline,
                                                     &down_args,
                                                     down_buf,
                                                     (NSUInteger)down_inner,
                                                     midbuf,
                                                     ds4_gpu_tensor_offset(mid),
                                                     outbuf,
                                                     ds4_gpu_tensor_offset(out),
                                                     selectedbuf,
                                                     ds4_gpu_tensor_offset(selected),
                                                     nil,
                                                     0,
                                                     down_smem,
                                                     2);
            } else if (use_mm_id) {
                ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
                                                       down_mm_pipeline,
                                                       &down_mm_args,
                                                       down_buf,
                                                       (NSUInteger)down_inner,
                                                       midbuf,
                                                       ds4_gpu_tensor_offset(mid),
                                                       down_dst,
                                                       down_dst_off,
                                                       8192u);
            } else {
                ok = ds4_gpu_encode_mul_mv_id(cb,
                                                     down_mv_pipeline,
                                                     &down_args,
                                                     down_buf,
                                                     (NSUInteger)down_inner,
                                                     midbuf,
                                                     ds4_gpu_tensor_offset(mid),
                                                     down_dst,
                                                     down_dst_off,
                                                     selectedbuf,
                                                     ds4_gpu_tensor_offset(selected),
                                                     down_smem,
                                                     down_nsg,
                                                     down_rows_per_group_is_nr0);
            }
        }
        DS4_METAL_PROFILE_MOE_STAGE("down");
        if (ok &&
            n_expert > 1 &&
            !direct_down_sum &&
            !use_q4_batch_expert_table &&
            !use_iq2_batch_selected_addr) {
            ok = ds4_gpu_encode_moe_sum_experts(cb,
                                                       down_dst,
                                                       down_dst_off,
                                                       outbuf,
                                                       ds4_gpu_tensor_offset(out),
                                                       out_dim,
                                                       n_expert,
                                                       n_tokens);
        }
        DS4_METAL_PROFILE_MOE_STAGE("sum");
        if (!ok) {
            fprintf(stderr,
                    "ds4: Metal routed batch MoE failed before submit layer=%u tokens=%u "
                    "gate=%s down=%s path=%s\n",
                    layer_index,
                    n_tokens,
                    ds4_gpu_metal_tensor_type_name(gate_type),
                    ds4_gpu_metal_tensor_type_name(down_type),
                    moe_path);
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "routed batch MoE")) {
            if (use_iq2_batch_selected_addr) {
                ds4_gpu_stream_expert_cache_clear_layer(layer_index);
            }
            return 0;
        }
        if (use_iq2_batch_selected_addr) {
            if (!owned) {
                if (ds4_gpu_end_commands() == 0) {
                    ds4_gpu_stream_expert_cache_clear_layer(layer_index);
                    return 0;
                }
                if (ds4_gpu_begin_commands() == 0) {
                    fprintf(stderr,
                            "ds4: Metal streaming prefill batch selected addr layer=%u "
                            "failed to reopen command batch after execution\n",
                            layer_index);
                    return 0;
                }
            }
        }
        if (q4_batch_table_boundary) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                return 0;
            }
        }
#undef DS4_METAL_PROFILE_MOE_STAGE
    }

    return 1;
}

int ds4_gpu_hc_split_sinkhorn_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *mix,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                scale_offset,
        uint64_t                base_offset,
        uint32_t                n_hc,
        uint32_t                sinkhorn_iters,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (n_hc == 0 || n_hc > 16) return 0;
    const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
    const uint64_t mix_bytes = mix_hc * sizeof(float);
    const uint64_t scale_bytes = 3ull * sizeof(float);

    @autoreleasepool {
        id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t mix_tensor_bytes = ds4_gpu_tensor_bytes(mix);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
        if (!mixbuf || !outbuf ||
            mix_tensor_bytes < mix_bytes ||
            out_tensor_bytes < mix_bytes) {
            fprintf(stderr, "ds4: Metal HC split received undersized activation buffers\n");
            return 0;
        }
        if (scale_offset > model_size || scale_bytes > model_size - scale_offset ||
            base_offset > model_size || mix_bytes > model_size - base_offset) {
            fprintf(stderr, "ds4: Metal HC split parameter range is outside the mapped model\n");
            return 0;
        }

        uint64_t scale_inner = 0;
        uint64_t base_inner = 0;
        id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner);
        id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
        if (!scalebuf || !basebuf) return 0;

        uint64_t n_rows64 = mix_tensor_bytes / mix_bytes;
        const uint64_t out_rows64 = out_tensor_bytes / mix_bytes;
        if (out_rows64 < n_rows64) n_rows64 = out_rows64;
        if (n_rows64 == 0 || n_rows64 > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal HC split row count is outside supported range\n");
            return 0;
        }

        ds4_gpu_hc_split_args args = {
            .n_hc = (int32_t)n_hc,
            .sinkhorn_iters = (int32_t)sinkhorn_iters,
            .n_rows = (int64_t)n_rows64,
            .mix_hc = (int64_t)mix_hc,
            .nb01 = mix_bytes,
            .nb1 = mix_bytes,
            .eps = eps,
        };
        const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_rows64));
        const NSUInteger n_tg = ((NSUInteger)n_rows64 + nth - 1u) / nth;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_hc_split_sinkhorn_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1];
        [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
        [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sinkhorn")) return 0;
    }

    return 1;
}

static int ds4_gpu_hc_weighted_sum_strided(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *weights,
        uint64_t                weight_offset,
        uint64_t                weight_row_stride,
        uint32_t                n_embd,
        uint32_t                n_hc,
        const char             *label) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0 ||
        weight_row_stride < (uint64_t)n_hc * sizeof(float)) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
        if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || out_tensor_bytes % out_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal HC weighted sum output size is not a whole token row\n");
            return 0;
        }

        const uint64_t n_tokens64 = out_tensor_bytes / out_row_bytes;
        if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal HC weighted sum token count is outside supported range\n");
            return 0;
        }

        const uint64_t x_row_values = (uint64_t)n_hc * n_embd;
        if (x_row_values == 0 ||
            x_row_values > UINT64_MAX / sizeof(float) ||
            n_tokens64 > UINT64_MAX / (x_row_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / ((uint64_t)n_hc * sizeof(float))) {
            fprintf(stderr, "ds4: Metal HC weighted sum activation size overflow\n");
            return 0;
        }

        const uint64_t x_bytes = n_tokens64 * x_row_values * sizeof(float);
        const uint64_t w_last = weight_offset +
                                (n_tokens64 - 1u) * weight_row_stride +
                                (uint64_t)n_hc * sizeof(float);
        if (!xbuf || !wbuf || !outbuf ||
            ds4_gpu_tensor_bytes(residual_hc) < x_bytes ||
            ds4_gpu_tensor_bytes(weights) < w_last) {
            fprintf(stderr, "ds4: Metal HC weighted sum received undersized activation buffers\n");
            return 0;
        }

        ds4_gpu_hc_weighted_sum_args args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = (int64_t)n_tokens64,
            .nb_x0 = sizeof(float),
            .nb_x1 = (uint64_t)n_embd * sizeof(float),
            .nb_x2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_w0 = sizeof(float),
            .nb_w1 = weight_row_stride,
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
        };
        const uint64_t n_elem = (uint64_t)n_embd * n_tokens64;
        const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
        const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_hc_weighted_sum_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1];
        [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) + (NSUInteger)weight_offset atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, label)) return 0;
    }

    return 1;
}

int ds4_gpu_hc_weighted_sum_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *weights,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    return ds4_gpu_hc_weighted_sum_strided(out,
                                             residual_hc,
                                             weights,
                                             0,
                                             (uint64_t)n_hc * sizeof(float),
                                             n_embd,
                                             n_hc,
                                             "HC weighted sum");
}

int ds4_gpu_hc_weighted_sum_norm_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *norm_out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *weights,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                norm_weight_offset,
        uint32_t                n_embd,
        uint32_t                n_hc,
        float                   norm_eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !norm_out || !residual_hc || !weights || !model_map) return 0;

    const bool force =
        getenv("DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") != NULL;
    const bool disabled =
        getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") != NULL;
    const bool require =
        getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL;
    const bool supported_shape =
        n_hc == 4u && (n_embd == 4096u || n_embd == 7168u);
    const bool auto_shape =
        n_embd == 4096u && ds4_gpu_device_name_contains("M3");
    const bool use_fusion =
        supported_shape && !g_quality_mode && !disabled &&
        g_hc_weighted_sum_norm_pipeline != nil &&
        (auto_shape || force);
    if (require && supported_shape && !use_fusion) {
        fprintf(stderr,
                "ds4: required Metal output HC sum/RMSNorm fusion was not selected\n");
        return 0;
    }
    if (!use_fusion) return 0;

    @autoreleasepool {
        const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
        const uint64_t residual_bytes = (uint64_t)n_hc * out_row_bytes;
        const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float);
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
        if (!xbuf || !wbuf || !outbuf || !normbuf ||
            ds4_gpu_tensor_bytes(residual_hc) < residual_bytes ||
            ds4_gpu_tensor_bytes(weights) < weight_bytes ||
            ds4_gpu_tensor_bytes(out) != out_row_bytes ||
            ds4_gpu_tensor_bytes(norm_out) < out_row_bytes) {
            fprintf(stderr,
                    "ds4: Metal output HC sum/RMSNorm fusion received invalid activation buffers\n");
            return 0;
        }
        if (norm_weight_offset > model_size ||
            out_row_bytes > model_size - norm_weight_offset) {
            fprintf(stderr,
                    "ds4: Metal output HC sum/RMSNorm weight range is outside the mapped model\n");
            return 0;
        }

        uint64_t norm_inner = 0;
        id<MTLBuffer> normwbuf = ds4_gpu_wrap_model_range(
            model_map, model_size, norm_weight_offset,
            out_row_bytes, &norm_inner);
        if (!normwbuf) return 0;

        ds4_gpu_hc_weighted_sum_norm_args args = {
            .n_embd = (int64_t)n_embd,
            .n_hc = (int64_t)n_hc,
            .n_tokens = 1,
            .nb_x0 = sizeof(float),
            .nb_x1 = out_row_bytes,
            .nb_x2 = residual_bytes,
            .nb_w0 = sizeof(float),
            .nb_w1 = weight_bytes,
            .nb0 = sizeof(float),
            .nb1 = out_row_bytes,
            .nb_norm1 = out_row_bytes,
            .norm_eps = norm_eps,
        };
        const NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd);
        const NSUInteger shared_bytes =
            ((NSUInteger)n_embd + 32u) * sizeof(float);
        if (nth > g_hc_weighted_sum_norm_pipeline.maxTotalThreadsPerThreadgroup ||
            shared_bytes > [g_device maxThreadgroupMemoryLength]) {
            if (require) {
                fprintf(stderr,
                        "ds4: required Metal output HC sum/RMSNorm fusion exceeds device limits\n");
            }
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_hc_weighted_sum_norm_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1];
        [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
        [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:4];
        [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5];
        [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(
                cb, owned, "output HC sum/RMSNorm fused")) {
            return 0;
        }
    }

    return 1;
}

int ds4_gpu_hc_weighted_sum_split_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
    return ds4_gpu_hc_weighted_sum_strided(out,
                                             residual_hc,
                                             split,
                                             0,
                                             mix_hc * sizeof(float),
                                             n_embd,
                                             n_hc,
                                             "HC weighted sum split");
}

/* Release decode fused HC pre-sublayer operation.  The graph driver owns the
 * optional reference fallback so this function stays a direct fused dispatch. */
int ds4_gpu_hc_split_weighted_sum_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *split,
        const ds4_gpu_tensor *mix,
        const ds4_gpu_tensor *residual_hc,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                scale_offset,
        uint64_t                base_offset,
        uint32_t                n_embd,
        uint32_t                n_hc,
        uint32_t                sinkhorn_iters,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !split || !mix || !residual_hc || !model_map ||
        n_embd == 0 || n_hc == 0) {
        return 0;
    }
    if (n_hc != 4) {
        fprintf(stderr, "ds4: Metal fused HC split/sum is specialized for HC=4\n");
        return 0;
    }

    const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
    const uint64_t mix_bytes = mix_hc * sizeof(float);
    const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
    const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
    const uint64_t scale_bytes = 3ull * sizeof(float);

    @autoreleasepool {
        id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
        id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
        if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes ||
            out_tensor_bytes % out_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal fused HC split/sum output size is not a whole token row\n");
            return 0;
        }

        const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes;
        if (n_rows64 == 0 || n_rows64 > UINT32_MAX ||
            n_rows64 > UINT64_MAX / mix_bytes ||
            n_rows64 > UINT64_MAX / residual_row_bytes) {
            fprintf(stderr, "ds4: Metal fused HC split/sum row count is outside supported range\n");
            return 0;
        }

        const uint64_t mix_total_bytes = n_rows64 * mix_bytes;
        const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes;
        if (!mixbuf || !splitbuf || !xbuf || !outbuf ||
            ds4_gpu_tensor_bytes(mix) < mix_total_bytes ||
            ds4_gpu_tensor_bytes(split) < mix_total_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes) {
            fprintf(stderr, "ds4: Metal fused HC split/sum received undersized activation buffers\n");
            return 0;
        }

        if (scale_offset > model_size || scale_bytes > model_size - scale_offset ||
            base_offset > model_size || mix_bytes > model_size - base_offset) {
            fprintf(stderr, "ds4: Metal fused HC split/sum parameter range is outside the mapped model\n");
            return 0;
        }

        uint64_t scale_inner = 0;
        uint64_t base_inner = 0;
        id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner);
        id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
        if (!scalebuf || !basebuf) return 0;

        ds4_gpu_hc_split_weighted_sum_args args = {
            .n_embd = (int64_t)n_embd,
            .n_hc = (int32_t)n_hc,
            .sinkhorn_iters = (int32_t)sinkhorn_iters,
            .n_rows = (int64_t)n_rows64,
            .mix_hc = (int64_t)mix_hc,
            .nb_mix1 = mix_bytes,
            .nb_split1 = mix_bytes,
            .nb_x0 = sizeof(float),
            .nb_x1 = (uint64_t)n_embd * sizeof(float),
            .nb_x2 = residual_row_bytes,
            .nb0 = sizeof(float),
            .nb1 = out_row_bytes,
            .eps = eps,
        };

        NSUInteger nth = g_hc_split_weighted_sum_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth > 256u) nth = 256u;
        if (nth > (NSUInteger)n_embd) nth = (NSUInteger)n_embd;
        if (nth == 0) nth = 1u;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_hc_split_weighted_sum_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1];
        [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
        [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6];
        [enc setThreadgroupMemoryLength:(NSUInteger)n_hc * sizeof(float) atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum fused")) return 0;
    }

    return 1;
}

/* HC-pre plus the immediately following weighted RMSNorm, specialized for
 * DS4's HC=4 shape. Both decode and batched prefill use this implementation;
 * the kernel preserves their established single-row and batched scale formulas. */
int ds4_gpu_hc_split_weighted_sum_norm_tensor(
        ds4_gpu_tensor       *out,
        ds4_gpu_tensor       *norm_out,
        ds4_gpu_tensor       *split,
        const ds4_gpu_tensor *mix,
        const ds4_gpu_tensor *residual_hc,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                scale_offset,
        uint64_t                base_offset,
        uint64_t                norm_weight_offset,
        uint32_t                n_embd,
        uint32_t                n_hc,
        uint32_t                sinkhorn_iters,
        float                   eps,
    float                   norm_eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !norm_out || !split || !mix || !residual_hc || !model_map ||
        n_hc != 4 || (n_embd & 3u) != 0) {
        return 0;
    }

    const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
    const uint64_t mix_bytes = mix_hc * sizeof(float);
    const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
    const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
    const uint64_t scale_bytes = 3ull * sizeof(float);

    @autoreleasepool {
        id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
        id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
        if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes ||
            out_tensor_bytes % out_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal fused HC split/sum/norm output size is not a whole token row\n");
            return 0;
        }

        const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes;
        if (n_rows64 == 0 || n_rows64 > UINT32_MAX ||
            n_rows64 > UINT64_MAX / mix_bytes ||
            n_rows64 > UINT64_MAX / residual_row_bytes) {
            fprintf(stderr, "ds4: Metal fused HC split/sum/norm row count is outside supported range\n");
            return 0;
        }

        const uint64_t mix_total_bytes = n_rows64 * mix_bytes;
        const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes;
        const uint64_t out_total_bytes = n_rows64 * out_row_bytes;
        if (!mixbuf || !splitbuf || !xbuf || !outbuf || !normbuf ||
            ds4_gpu_tensor_bytes(mix) < mix_total_bytes ||
            ds4_gpu_tensor_bytes(split) < mix_total_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes ||
            ds4_gpu_tensor_bytes(norm_out) < out_total_bytes) {
            fprintf(stderr, "ds4: Metal fused HC split/sum/norm received undersized activation buffers\n");
            return 0;
        }

        if (scale_offset > model_size || scale_bytes > model_size - scale_offset ||
            base_offset > model_size || mix_bytes > model_size - base_offset ||
            norm_weight_offset > model_size || out_row_bytes > model_size - norm_weight_offset) {
            fprintf(stderr, "ds4: Metal fused HC split/sum/norm parameter range is outside the mapped model\n");
            return 0;
        }

        uint64_t scale_inner = 0;
        uint64_t base_inner = 0;
        uint64_t norm_inner = 0;
        id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner);
        id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
        id<MTLBuffer> normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, out_row_bytes, &norm_inner);
        if (!scalebuf || !basebuf || !normwbuf) return 0;

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_hot_pipeline(g_hc_split_weighted_sum_norm_pipeline,
                                   "kernel_dsv4_hc_split_weighted_sum_norm4");
        if (!pipeline) return 0;

        ds4_gpu_hc_split_weighted_sum_norm_args args = {
            .n_embd = (int64_t)n_embd,
            .n_hc = (int32_t)n_hc,
            .sinkhorn_iters = (int32_t)sinkhorn_iters,
            .n_rows = (int64_t)n_rows64,
            .mix_hc = (int64_t)mix_hc,
            .nb_mix1 = mix_bytes,
            .nb_split1 = mix_bytes,
            .nb_x0 = sizeof(float),
            .nb_x1 = (uint64_t)n_embd * sizeof(float),
            .nb_x2 = residual_row_bytes,
            .nb0 = sizeof(float),
            .nb1 = out_row_bytes,
            .nb_norm1 = out_row_bytes,
            .eps = eps,
            .norm_eps = norm_eps,
        };

        NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd);
        if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
            fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu threads but pipeline supports %lu\n",
                    (unsigned long)nth,
                    (unsigned long)pipeline.maxTotalThreadsPerThreadgroup);
            return 0;
        }

        const NSUInteger shared_bytes = ((NSUInteger)n_embd + 4u + 32u) * sizeof(float);
        const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength];
        if (max_shared != 0 && shared_bytes > max_shared) {
            fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu bytes of threadgroup memory but device supports %lu\n",
                    (unsigned long)shared_bytes,
                    (unsigned long)max_shared);
            return 0;
        }

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1];
        [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
        [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6];
        [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:7];
        [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:8];
        [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum/norm fused")) return 0;
    }

    return 1;
}

int ds4_gpu_output_hc_weights_tensor(
        ds4_gpu_tensor       *out,
        const ds4_gpu_tensor *pre,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                scale_offset,
        uint64_t                base_offset,
        uint32_t                n_hc,
        float                   eps) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !pre || !model_map || n_hc == 0) return 0;

    @autoreleasepool {
        if ((n_hc % 4u) != 0) {
            fprintf(stderr, "ds4: Metal output HC weights requires a multiple-of-4 HC width\n");
            return 0;
        }

        id<MTLBuffer> prebuf = ds4_gpu_tensor_buffer(pre);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
        if (row_bytes == 0 || out_tensor_bytes < row_bytes || out_tensor_bytes % row_bytes != 0) {
            fprintf(stderr, "ds4: Metal output HC weights size is not a whole token row\n");
            return 0;
        }

        const uint64_t n_tokens64 = out_tensor_bytes / row_bytes;
        if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX ||
            n_tokens64 > UINT64_MAX / row_bytes) {
            fprintf(stderr, "ds4: Metal output HC weights token count is outside supported range\n");
            return 0;
        }

        const uint64_t bytes = n_tokens64 * row_bytes;
        if (!prebuf || !outbuf ||
            ds4_gpu_tensor_bytes(pre) < bytes ||
            ds4_gpu_tensor_bytes(out) < bytes) {
            fprintf(stderr, "ds4: Metal output HC weights received undersized buffers\n");
            return 0;
        }

        uint64_t scale_inner = 0;
        uint64_t base_inner = 0;
        id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                            scale_offset, sizeof(float),
                                                            &scale_inner);
        id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                           base_offset, row_bytes,
                                                           &base_inner);
        if (!scalebuf || !basebuf) return 0;

        const bool force_weights4 =
            getenv("DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4") != NULL;
        const bool disable_weights4 =
            getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4") != NULL;
        const bool require_weights4 =
            getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL;
        const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u;
        const bool use_weights4 =
            weights4_shape && !g_quality_mode && !disable_weights4 &&
            g_output_hc_weights4_pipeline != nil &&
            (ds4_gpu_device_name_contains("M3") || force_weights4) &&
            g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u;
        if (require_weights4 && weights4_shape && !use_weights4) {
            fprintf(stderr,
                    "ds4: required Metal output HC weights4 kernel was not selected\n");
            return 0;
        }

        if (use_weights4) {
            ds4_gpu_output_hc_weights4_args args = {
                .post_scale = 1.0f,
                .eps = eps,
            };
            int owned = 0;
            id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
            if (!cb) return 0;

            id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
            [enc setComputePipelineState:g_output_hc_weights4_pipeline];
            [enc setBytes:&args length:sizeof(args) atIndex:0];
            [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1];
            [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
            [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
            [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
            [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
                 threadsPerThreadgroup:MTLSizeMake(2, 1, 1)];
            ds4_gpu_end_compute_encoder(cb, enc);

            if (!ds4_gpu_finish_command_buffer(
                    cb, owned, "output HC weights4")) {
                return 0;
            }
            return 1;
        }

        const uint32_t n_tokens = (uint32_t)n_tokens64;
        ds4_gpu_bin_args mul_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, 1);
        ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, n_hc);
        ds4_gpu_unary_args sigmoid_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 0.0f, 0.0f);
        ds4_gpu_unary_args scale_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 1.0f, eps);

        NSUInteger mul_nth_max = g_bin_mul_scalar_pipeline.maxTotalThreadsPerThreadgroup;
        if (mul_nth_max > 256u) mul_nth_max = 256u;
        NSUInteger mul_nth = 1u;
        while (2u * mul_nth < (NSUInteger)mul_args.ne0 && mul_nth < mul_nth_max) {
            mul_nth *= 2u;
        }

        NSUInteger add_nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup;
        if (add_nth_max > 256u) add_nth_max = 256u;
        NSUInteger add_nth = 1u;
        while (2u * add_nth < (NSUInteger)add_args.ne0 && add_nth < add_nth_max) {
            add_nth *= 2u;
        }

        NSUInteger unary_nth_max = g_unary_sigmoid_pipeline.maxTotalThreadsPerThreadgroup;
        if (unary_nth_max > 256u) unary_nth_max = 256u;
        NSUInteger unary_nth = (NSUInteger)sigmoid_args.ne00;
        if (unary_nth > unary_nth_max) unary_nth = unary_nth_max;
        if (unary_nth == 0) unary_nth = 1u;
        const NSUInteger unary_nk0 = ((NSUInteger)sigmoid_args.ne00 + unary_nth - 1u) / unary_nth;
        const NSUInteger out_offset = ds4_gpu_tensor_offset(out);

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);

        [enc setComputePipelineState:g_bin_mul_scalar_pipeline];
        [enc setBytes:&mul_args length:sizeof(mul_args) atIndex:0];
        [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1];
        [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
        [enc setBuffer:outbuf offset:out_offset atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)mul_args.ne01,
                                              (NSUInteger)mul_args.ne02,
                                              (NSUInteger)mul_args.ne03)
             threadsPerThreadgroup:MTLSizeMake(mul_nth, 1, 1)];

        [enc setComputePipelineState:g_add_pipeline];
        [enc setBytes:&add_args length:sizeof(add_args) atIndex:0];
        [enc setBuffer:outbuf offset:out_offset atIndex:1];
        [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:2];
        [enc setBuffer:outbuf offset:out_offset atIndex:3];
        [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)add_args.ne01,
                                              (NSUInteger)add_args.ne02,
                                              (NSUInteger)add_args.ne03)
             threadsPerThreadgroup:MTLSizeMake(add_nth, 1, 1)];

        [enc setComputePipelineState:g_unary_sigmoid_pipeline];
        [enc setBytes:&sigmoid_args length:sizeof(sigmoid_args) atIndex:0];
        [enc setBuffer:outbuf offset:out_offset atIndex:1];
        [enc setBuffer:outbuf offset:out_offset atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)sigmoid_args.ne01,
                                              (NSUInteger)sigmoid_args.ne02,
                                              (NSUInteger)sigmoid_args.ne03)
             threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)];

        [enc setComputePipelineState:g_unary_scale_pipeline];
        [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0];
        [enc setBuffer:outbuf offset:out_offset atIndex:1];
        [enc setBuffer:outbuf offset:out_offset atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)scale_args.ne01,
                                              (NSUInteger)scale_args.ne02,
                                              (NSUInteger)scale_args.ne03)
             threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "output HC weights")) return 0;
    }

    return 1;
}

int ds4_gpu_hc_expand_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *block_out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *post,
        const ds4_gpu_tensor *comb,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (n_embd == 0 || n_hc == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
        id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> postbuf = ds4_gpu_tensor_buffer(post);
        id<MTLBuffer> combbuf = ds4_gpu_tensor_buffer(comb);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
        const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
        if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n");
            return 0;
        }

        const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
        if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n");
            return 0;
        }

        const uint64_t block_values = (uint64_t)n_embd;
        const uint64_t hc_values = (uint64_t)n_hc * n_embd;
        const uint64_t comb_values = (uint64_t)n_hc * n_hc;
        if (hc_values == 0 ||
            hc_values > UINT64_MAX / sizeof(float) ||
            comb_values > UINT64_MAX / sizeof(float) ||
            n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) {
            fprintf(stderr, "ds4: Metal HC expand activation size overflow\n");
            return 0;
        }

        const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
        const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
        const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float);
        const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float);
        if (!blockbuf || !resbuf || !postbuf || !combbuf || !outbuf ||
            ds4_gpu_tensor_bytes(block_out) < block_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
            ds4_gpu_tensor_bytes(post) < post_bytes ||
            ds4_gpu_tensor_bytes(comb) < comb_bytes) {
            fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n");
            return 0;
        }

        ds4_gpu_hc_expand_args args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = (int64_t)n_tokens64,
            .nb_block0 = sizeof(float),
            .nb_block1 = (uint64_t)n_embd * sizeof(float),
            .nb_add0 = sizeof(float),
            .nb_add1 = (uint64_t)n_embd * sizeof(float),
            .nb_res0 = sizeof(float),
            .nb_res1 = (uint64_t)n_embd * sizeof(float),
            .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_post0 = sizeof(float),
            .nb_post1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb0 = sizeof(float),
            .nb_comb1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
            .nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .has_add = 0,
        };
        id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
        uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
        if (n_hc == 4) {
            expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
                                                      "kernel_dsv4_hc_expand4");
            n_elem = (uint64_t)n_embd * n_tokens64;
        }
        if (!expand_pipeline) return 0;
        const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
        const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:expand_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
        [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
        [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3];
        [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand")) return 0;
    }

    return 1;
}

/* Expand of the SUM of two block vectors — the TP attention combine folded
 * into the HC expand (the kernel's has_add path adds them element-wise
 * before the post/comb mixing, canonical rank order preserved by argument
 * position). */
int ds4_gpu_hc_expand_add_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *block_out,
        const ds4_gpu_tensor *block_add,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *post,
        const ds4_gpu_tensor *comb,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (n_embd == 0 || n_hc != 4) return 0;

    @autoreleasepool {
        id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
        id<MTLBuffer> addbuf = ds4_gpu_tensor_buffer(block_add);
        id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> postbuf = ds4_gpu_tensor_buffer(post);
        id<MTLBuffer> combbuf = ds4_gpu_tensor_buffer(comb);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
        const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
        if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n");
            return 0;
        }

        const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
        if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n");
            return 0;
        }

        const uint64_t block_values = (uint64_t)n_embd;
        const uint64_t hc_values = (uint64_t)n_hc * n_embd;
        const uint64_t comb_values = (uint64_t)n_hc * n_hc;
        if (hc_values == 0 ||
            hc_values > UINT64_MAX / sizeof(float) ||
            comb_values > UINT64_MAX / sizeof(float) ||
            n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) {
            fprintf(stderr, "ds4: Metal HC expand activation size overflow\n");
            return 0;
        }

        const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
        const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
        const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float);
        const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float);
        if (!blockbuf || !addbuf || !resbuf || !postbuf || !combbuf || !outbuf ||
            ds4_gpu_tensor_bytes(block_out) < block_bytes ||
            ds4_gpu_tensor_bytes(block_add) < block_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
            ds4_gpu_tensor_bytes(post) < post_bytes ||
            ds4_gpu_tensor_bytes(comb) < comb_bytes) {
            fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n");
            return 0;
        }

        ds4_gpu_hc_expand_args args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = (int64_t)n_tokens64,
            .nb_block0 = sizeof(float),
            .nb_block1 = (uint64_t)n_embd * sizeof(float),
            .nb_add0 = sizeof(float),
            .nb_add1 = (uint64_t)n_embd * sizeof(float),
            .nb_res0 = sizeof(float),
            .nb_res1 = (uint64_t)n_embd * sizeof(float),
            .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_post0 = sizeof(float),
            .nb_post1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb0 = sizeof(float),
            .nb_comb1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
            .nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .has_add = 1,
        };
        id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
        uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
        if (n_hc == 4) {
            expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
                                                      "kernel_dsv4_hc_expand4");
            n_elem = (uint64_t)n_embd * n_tokens64;
        }
        if (!expand_pipeline) return 0;
        const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
        const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:expand_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
        [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
        [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3];
        [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4];
        [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add")) return 0;
    }

    return 1;
}

int ds4_gpu_hc_expand_split_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *block_out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
        id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
        const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
        if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal HC expand split output size is not a whole HC token row\n");
            return 0;
        }

        const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
        if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal HC expand split token count is outside supported range\n");
            return 0;
        }

        const uint64_t block_values = (uint64_t)n_embd;
        const uint64_t hc_values = (uint64_t)n_hc * n_embd;
        const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
        if (hc_values == 0 ||
            hc_values > UINT64_MAX / sizeof(float) ||
            mix_hc > UINT64_MAX / sizeof(float) ||
            n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) {
            fprintf(stderr, "ds4: Metal HC expand split activation size overflow\n");
            return 0;
        }

        const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
        const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
        const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float);
        if (!blockbuf || !resbuf || !splitbuf || !outbuf ||
            ds4_gpu_tensor_bytes(block_out) < block_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
            ds4_gpu_tensor_bytes(split) < split_bytes) {
            fprintf(stderr, "ds4: Metal HC expand split received undersized activation buffers\n");
            return 0;
        }

        ds4_gpu_hc_expand_args args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = (int64_t)n_tokens64,
            .nb_block0 = sizeof(float),
            .nb_block1 = (uint64_t)n_embd * sizeof(float),
            .nb_add0 = sizeof(float),
            .nb_add1 = (uint64_t)n_embd * sizeof(float),
            .nb_res0 = sizeof(float),
            .nb_res1 = (uint64_t)n_embd * sizeof(float),
            .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_post0 = sizeof(float),
            .nb_post1 = mix_hc * sizeof(float),
            .nb_comb0 = sizeof(float),
            .nb_comb1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb2 = mix_hc * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
            .nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .has_add = 0,
        };
        id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
        uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
        if (n_hc == 4) {
            expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
                                                      "kernel_dsv4_hc_expand4");
            n_elem = (uint64_t)n_embd * n_tokens64;
        }
        if (!expand_pipeline) return 0;
        const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
        const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:expand_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
        [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand split")) return 0;
    }

    return 1;
}

int ds4_gpu_hc_expand_split_half_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *block_out_h,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    (void)out_hc; (void)block_out_h; (void)residual_hc; (void)split;
    (void)n_embd; (void)n_hc;
    return 0;
}

int ds4_gpu_hc_expand_add_split_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *block_out,
        const ds4_gpu_tensor *block_add,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;

    @autoreleasepool {
        id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
        id<MTLBuffer> addbuf = ds4_gpu_tensor_buffer(block_add);
        id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
        const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
        const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
        if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
            fprintf(stderr, "ds4: Metal HC expand add split output size is not a whole HC token row\n");
            return 0;
        }

        const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
        if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
            fprintf(stderr, "ds4: Metal HC expand add split token count is outside supported range\n");
            return 0;
        }

        const uint64_t block_values = (uint64_t)n_embd;
        const uint64_t hc_values = (uint64_t)n_hc * n_embd;
        const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
        if (hc_values == 0 ||
            hc_values > UINT64_MAX / sizeof(float) ||
            mix_hc > UINT64_MAX / sizeof(float) ||
            n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
            n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) {
            fprintf(stderr, "ds4: Metal HC expand add split activation size overflow\n");
            return 0;
        }

        const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
        const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
        const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float);
        if (!blockbuf || !addbuf || !resbuf || !splitbuf || !outbuf ||
            ds4_gpu_tensor_bytes(block_out) < block_bytes ||
            ds4_gpu_tensor_bytes(block_add) < block_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
            ds4_gpu_tensor_bytes(split) < split_bytes) {
            fprintf(stderr, "ds4: Metal HC expand add split received undersized activation buffers\n");
            return 0;
        }

        ds4_gpu_hc_expand_args args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = (int64_t)n_tokens64,
            .nb_block0 = sizeof(float),
            .nb_block1 = (uint64_t)n_embd * sizeof(float),
            .nb_add0 = sizeof(float),
            .nb_add1 = (uint64_t)n_embd * sizeof(float),
            .nb_res0 = sizeof(float),
            .nb_res1 = (uint64_t)n_embd * sizeof(float),
            .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_post0 = sizeof(float),
            .nb_post1 = mix_hc * sizeof(float),
            .nb_comb0 = sizeof(float),
            .nb_comb1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb2 = mix_hc * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
            .nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .has_add = 1,
        };
        id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
        uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
        if (n_hc == 4) {
            expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
                                                      "kernel_dsv4_hc_expand4");
            n_elem = (uint64_t)n_embd * n_tokens64;
        }
        if (!expand_pipeline) return 0;
        const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
        const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:expand_pipeline];
        [enc setBytes:&args length:sizeof(args) atIndex:0];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
        [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4];
        [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
             threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add split")) return 0;
    }

    return 1;
}

int ds4_gpu_hc_expand_add_split_half_add_tensor(
        ds4_gpu_tensor       *out_hc,
        const ds4_gpu_tensor *block_out,
        const ds4_gpu_tensor *block_add_h,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    (void)out_hc; (void)block_out; (void)block_add_h; (void)residual_hc;
    (void)split; (void)n_embd; (void)n_hc;
    return 0;
}

int ds4_gpu_shared_down_hc_expand_q8_0_tensor(
        ds4_gpu_tensor       *out_hc,
        ds4_gpu_tensor       *shared_out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *shared_mid,
        const ds4_gpu_tensor *routed_out,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out_hc || !shared_out || !model_map || !shared_mid || !routed_out ||
        !residual_hc || !split || n_embd == 0 || n_hc == 0 ||
        n_hc != 4 || out_dim != n_embd || (in_dim & 31u) != 0 ||
        in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(shared_mid);
        id<MTLBuffer> sharedbuf = ds4_gpu_tensor_buffer(shared_out);
        id<MTLBuffer> routedbuf = ds4_gpu_tensor_buffer(routed_out);
        id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);

        const uint64_t row_bytes = (in_dim / 32u) * 34u;
        const uint64_t weight_bytes = out_dim * row_bytes;
        const uint64_t shared_mid_bytes = in_dim * sizeof(float);
        const uint64_t embd_bytes = out_dim * sizeof(float);
        const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
        const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
        const uint64_t split_bytes = mix_hc * sizeof(float);

        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal shared-down HC fusion weight range is outside the mapped model\n");
            return 0;
        }
        if (!midbuf || !sharedbuf || !routedbuf || !resbuf || !splitbuf || !outbuf ||
            ds4_gpu_tensor_bytes(shared_mid) < shared_mid_bytes ||
            ds4_gpu_tensor_bytes(shared_out) < embd_bytes ||
            ds4_gpu_tensor_bytes(routed_out) < embd_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
            ds4_gpu_tensor_bytes(split) < split_bytes ||
            ds4_gpu_tensor_bytes(out_hc) < hc_bytes) {
            fprintf(stderr, "ds4: Metal shared-down HC fusion received undersized buffers\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                        weight_offset, weight_bytes,
                                                        &inner_offset);
        if (!wbuf) return 0;

        ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        mv_args.nr0 = mv_dispatch.nr0;

        ds4_gpu_hc_expand_args hc_args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = 1,
            .nb_block0 = sizeof(float),
            .nb_block1 = (uint64_t)n_embd * sizeof(float),
            .nb_add0 = sizeof(float),
            .nb_add1 = (uint64_t)n_embd * sizeof(float),
            .nb_res0 = sizeof(float),
            .nb_res1 = (uint64_t)n_embd * sizeof(float),
            .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_post0 = sizeof(float),
            .nb_post1 = mix_hc * sizeof(float),
            .nb_comb0 = sizeof(float),
            .nb_comb1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb2 = mix_hc * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
            .nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .has_add = 1,
        };

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_down_hc_expand4_q8_0",
                                          mv_dispatch.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
        [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2];
        [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(shared_mid) atIndex:3];
        [enc setBuffer:sharedbuf offset:ds4_gpu_tensor_offset(shared_out) atIndex:4];
        [enc setBuffer:routedbuf offset:ds4_gpu_tensor_offset(routed_out) atIndex:5];
        [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:6];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:7];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:8];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:9];
        [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
                                              (NSUInteger)mv_dispatch.nr0,
                                              1,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "shared-down HC expand fused")) return 0;
    }

    return 1;
}

int ds4_gpu_matmul_q8_0_hc_expand_tensor(
        ds4_gpu_tensor       *out_hc,
        ds4_gpu_tensor       *block_out,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                weight_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        const ds4_gpu_tensor *residual_hc,
        const ds4_gpu_tensor *split,
        uint32_t                n_embd,
        uint32_t                n_hc) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out_hc || !block_out || !model_map || !x || !residual_hc || !split ||
        n_embd == 0 || n_hc == 0 || n_hc != 4 || out_dim != n_embd ||
        (in_dim & 31u) != 0 || in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
        return 0;
    }

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
        id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
        id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);

        const uint64_t row_bytes = (in_dim / 32u) * 34u;
        const uint64_t weight_bytes = out_dim * row_bytes;
        const uint64_t x_bytes = in_dim * sizeof(float);
        const uint64_t embd_bytes = out_dim * sizeof(float);
        const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
        const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
        const uint64_t split_bytes = mix_hc * sizeof(float);

        if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
            fprintf(stderr, "ds4: Metal Q8 HC fusion weight range is outside the mapped model\n");
            return 0;
        }
        if (!xbuf || !blockbuf || !resbuf || !splitbuf || !outbuf ||
            ds4_gpu_tensor_bytes(x) < x_bytes ||
            ds4_gpu_tensor_bytes(block_out) < embd_bytes ||
            ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
            ds4_gpu_tensor_bytes(split) < split_bytes ||
            ds4_gpu_tensor_bytes(out_hc) < hc_bytes) {
            fprintf(stderr, "ds4: Metal Q8 HC fusion received undersized buffers\n");
            return 0;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
                                                        weight_offset, weight_bytes,
                                                        &inner_offset);
        if (!wbuf) return 0;

        ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
        ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
        mv_args.nr0 = mv_dispatch.nr0;

        ds4_gpu_hc_expand_args hc_args = {
            .n_embd = n_embd,
            .n_hc = n_hc,
            .n_tokens = 1,
            .nb_block0 = sizeof(float),
            .nb_block1 = (uint64_t)n_embd * sizeof(float),
            .nb_add0 = sizeof(float),
            .nb_add1 = (uint64_t)n_embd * sizeof(float),
            .nb_res0 = sizeof(float),
            .nb_res1 = (uint64_t)n_embd * sizeof(float),
            .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .nb_post0 = sizeof(float),
            .nb_post1 = mix_hc * sizeof(float),
            .nb_comb0 = sizeof(float),
            .nb_comb1 = (uint64_t)n_hc * sizeof(float),
            .nb_comb2 = mix_hc * sizeof(float),
            .nb0 = sizeof(float),
            .nb1 = (uint64_t)n_embd * sizeof(float),
            .nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
            .has_add = 0,
        };

        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_q8_hc_expand4_q8_0",
                                          mv_dispatch.nsg);
        if (!pipeline) return 0;

        int owned = 0;
        id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
        if (!cb) return 0;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:pipeline];
        [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
        [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1];
        [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2];
        [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
        [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:4];
        [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:5];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:6];
        [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7];
        [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:8];
        [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
        [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
                                              (NSUInteger)mv_dispatch.nr0,
                                              1,
                                              1)
             threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
        ds4_gpu_end_compute_encoder(cb, enc);

        if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8 HC expand fused")) return 0;
    }

    return 1;
}

void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) {
    (void)enabled;
}