antirez--ds4/ds4_metal.m

#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 C engine.
 *
 * The C code 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 C while still using Objective-C where Metal
 * requires it.
 */

enum {
    DS4_METAL_TENSOR_Q2_K    = 10,
    DS4_METAL_TENSOR_Q4_K    = 12,
    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 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_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_f16_f32_pipeline;
static id<MTLComputePipelineState> g_swiglu_pipeline;
static id<MTLComputePipelineState> g_add_pipeline;
static id<MTLComputePipelineState> g_moe_sum6_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_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_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_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_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_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_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_weights_one_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;
static id<MTLBuffer> g_flash_attn_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_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 uint32_t g_stream_expert_cache_mlock_budget_cap;
static uint8_t g_stream_expert_cache_mlock_relief_applied;
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_mlock_bytes;
static uint64_t g_stream_expert_cache_mlock_fail_bytes;
static uint64_t g_stream_expert_cache_mlock_failures;
static double g_stream_expert_cache_mlock_ms;
static int g_stream_expert_cache_mlock_warned;
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_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_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_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 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 uint64_t g_model_residency_count;
static int g_model_residency_added_to_queue;
static int g_ssd_streaming_mode;
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);
static NSUInteger g_flash_attn_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_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;
static int32_t g_routed_moe_selected_override[6];
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 = 61,
    DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384,
    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 uint8_t g_stream_expert_cache_slab_slot_locked[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES];
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) {
        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 int ds4_gpu_wait_command_buffer(id<MTLCommandBuffer> cb, const char *label) {
    [cb waitUntilCompleted];
    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();
    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_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 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[6],
        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 > 6) 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;
}

static int ds4_gpu_model_residency_request_views(void) {
    if (g_model_view_count == 0 ||
        g_ssd_streaming_mode ||
        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_USE_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.
         */
        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 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 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,
        int32_t     ns10,
        int32_t     ns20,
        int32_t     nsg,
        int32_t     nwg) {
    NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%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,
                     (int)ns10,
                     (int)ns20,
                     (int)nsg,
                     (int)nwg];
    id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
    if (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:&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];
    return pipeline;
}

static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_reduce_pipeline(
        int32_t dv,
        int32_t nwg) {
    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) 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];
    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_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,
        .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_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),
        .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_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 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 bind_avg=%.3f ms read_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 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_bind_avg,
            s.selected_read_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.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) {
    const uint64_t scratch =
        (uint64_t)g_flash_attn_mask_bytes +
        (uint64_t)g_flash_attn_pad_bytes +
        (uint64_t)g_flash_attn_tmp_bytes +
        (uint64_t)g_flash_attn_blk_bytes +
        (uint64_t)g_flash_attn_ring_bytes +
        (uint64_t)g_flash_attn_kv_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;

    fprintf(stderr, "ds4: Metal memory report%s%s\n",
            label && label[0] ? " " : "",
            label && label[0] ? label : "");
    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);
    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 (g_stream_expert_cache_mlock_bytes != 0 ||
            g_stream_expert_cache_mlock_failures != 0) {
            fprintf(stderr,
                    "ds4:   streaming expert buffer mlock locked=%.2f GiB failed=%.2f GiB failures=%llu time=%.3f ms\n",
                    ds4_gpu_gib(g_stream_expert_cache_mlock_bytes),
                    ds4_gpu_gib(g_stream_expert_cache_mlock_fail_bytes),
                    (unsigned long long)g_stream_expert_cache_mlock_failures,
                    g_stream_expert_cache_mlock_ms);
        }
        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:   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),
            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),
            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));
}

void ds4_gpu_set_quality(bool quality) {
    g_quality_mode = quality ? 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_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 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"
"#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"
"\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  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 {
    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 ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) {
    return (ds4_gpu_mv_dispatch) {
        .function_name = "kernel_mul_mv_q8_0_f32",
        .nsg = 4,
        .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 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;
    }
}

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;
    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;
} 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;
} 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);

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);

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 {
    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;

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) {
    int32_t pos_stack[256];
    int32_t *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);

    const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(*pos);
    id<MTLBuffer> posbuf = nil;
    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 instead of reporting a normal
         * API error.
         */
        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);
    }

    const NSUInteger nth = (NSUInteger)(head_dim < 256u ? head_dim : 256u);
    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_rope_tail_batch_pipeline];
    [enc setBytes:args length:sizeof(*args) atIndex:0];
    [enc setBuffer:xbuf offset:xoff atIndex:1];
    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(n_head, n_tok, 1)
         threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    if (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 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 {
    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_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_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_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_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_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;
        }

        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;
        }

        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_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_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_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;
        }

        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;
        }

        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_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_weights_one_pipeline =
            ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one");
        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_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();
    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_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();
    if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
    return g_batch_cb != nil;
}

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;
    [cb commit];
    [g_pending_cbs addObject:cb];
    ds4_gpu_stream_expert_cache_note_batch_committed();

    g_batch_cb = ds4_gpu_new_command_buffer();
    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;
}

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];
        *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;
        [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();
        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;
}

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];
        [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();
        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;
        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_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");
        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_repeat_f32_pipeline = nil;
        g_concat_pipeline = nil;
        g_cpy_f32_f32_pipeline = nil;
        g_cpy_f32_f16_pipeline = nil;
        g_cpy_f16_f32_pipeline = nil;
        g_swiglu_pipeline = nil;
        g_add_pipeline = nil;
        g_moe_sum6_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_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_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_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_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_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_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_weights_one_pipeline = nil;
        g_dsv4_hc_expand4_pipeline = nil;
        g_flash_attn_mask_buffer = nil;
        g_flash_attn_pad_buffer = nil;
        g_flash_attn_tmp_buffer = nil;
        g_flash_attn_blk_buffer = nil;
        g_flash_attn_ring_buffer = nil;
        g_flash_attn_kv_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_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_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 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_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);
}

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 slot budget and deadlock slab reuse after an mlock cap.
     */
    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) {
    uint32_t budget = ds4_gpu_stream_expert_cache_requested_budget();
    if (budget != 0 &&
        g_stream_expert_cache_mlock_budget_cap != 0 &&
        budget > g_stream_expert_cache_mlock_budget_cap) {
        budget = g_stream_expert_cache_mlock_budget_cap;
    }
    return 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_popcount6(uint32_t mask) {
    return (uint32_t)__builtin_popcount(mask & 0x3fu);
}

static void ds4_gpu_stream_expert_timing_note_selected(
        double read_ms,
        double bind_ms) {
    if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
    g_stream_expert_timing_selected_calls++;
    g_stream_expert_timing_selected_read_ms += read_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_popcount6(resident_mask);
    g_stream_expert_timing_split_missing_experts +=
        ds4_gpu_stream_expert_popcount6(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_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_popcount6(resident_mask);
    const uint32_t missing = ds4_gpu_stream_expert_popcount6(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)
    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;
    }
#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[6];
    uint32_t source_slots[6];
    uint32_t n_loads;
    uint32_t n_tasks;
    uint64_t gate_expert_bytes;
    uint64_t down_expert_bytes;
    int32_t selected_ids[6];
    uint64_t gate_abs_offsets[6];
    uint64_t up_abs_offsets[6];
    uint64_t down_abs_offsets[6];
    __strong id<MTLBuffer> gate_bufs[6];
    __strong id<MTLBuffer> up_bufs[6];
    __strong id<MTLBuffer> down_bufs[6];
    NSUInteger gate_inners[6];
    NSUInteger up_inners[6];
    NSUInteger down_inners[6];
    ds4_gpu_stream_expert_pread_task tasks[18];
    double start_ms;
} ds4_gpu_stream_expert_pending_load;

static ds4_gpu_stream_expert_pending_load g_stream_expert_pending_load;

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 void ds4_gpu_stream_expert_cache_warn_mlock_failure(
        uint64_t failed_len,
        int      err) {
    if (g_stream_expert_cache_mlock_warned) return;
    g_stream_expert_cache_mlock_warned = 1;

    const uint64_t gib = 1024ull * 1024ull * 1024ull;
    const uint32_t budget = ds4_gpu_stream_expert_cache_requested_budget();
    uint64_t requested = 0;
    if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) {
        requested =
            budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ?
                UINT64_MAX :
                (uint64_t)budget * g_stream_expert_cache_expert_bytes;
    }

    uint64_t suggested_gib = g_stream_expert_cache_mlock_bytes / gib;
    if (suggested_gib > 1) suggested_gib--;

    fprintf(stderr,
            "ds4: warning: streaming expert cache could not mlock all buffers\n");
    if (requested != 0) {
        fprintf(stderr,
                "ds4:   requested cache: %u experts / %.2f GiB\n",
                budget,
                ds4_gpu_gib(requested));
    } else {
        fprintf(stderr,
                "ds4:   requested cache: %u experts\n",
                budget);
    }
    fprintf(stderr,
            "ds4:   locked so far:   %.2f GiB\n",
            ds4_gpu_gib(g_stream_expert_cache_mlock_bytes));
    fprintf(stderr,
            "ds4:   failed buffer:   %.2f MiB (%s)\n",
            ds4_gpu_mib(failed_len),
            err != 0 ? strerror(err) : "mlock unavailable");
    fprintf(stderr,
            "ds4:   macOS may page unlocked expert buffers, causing poor or unstable speed\n");
    if (g_stream_expert_cache_mlock_budget_cap != 0 &&
        g_stream_expert_cache_expert_bytes != 0) {
        const uint64_t capped_bytes =
            g_stream_expert_cache_mlock_budget_cap >
                UINT64_MAX / g_stream_expert_cache_expert_bytes ?
                    UINT64_MAX :
                    (uint64_t)g_stream_expert_cache_mlock_budget_cap *
                        g_stream_expert_cache_expert_bytes;
        fprintf(stderr,
                "ds4:   using locked cache cap: %u experts / %.2f GiB\n",
                g_stream_expert_cache_mlock_budget_cap,
                ds4_gpu_gib(capped_bytes));
    }
    if (suggested_gib != 0) {
        fprintf(stderr,
                "ds4:   try: --ssd-streaming-cache-experts %" PRIu64 "GB\n",
                suggested_gib);
    } else {
        fprintf(stderr,
                "ds4:   try a smaller --ssd-streaming-cache-experts NGB budget\n");
    }
}

static uint32_t ds4_gpu_stream_expert_slab_locked_slot_count(void) {
    uint32_t count = 0;
    uint32_t total = g_stream_expert_cache_slab_total_slots;
    if (total > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
        total = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES;
    }
    for (uint32_t slot = 0; slot < total; slot++) {
        if (g_stream_expert_cache_slab_slot_locked[slot]) count++;
    }
    return count;
}

static void ds4_gpu_stream_expert_cache_cap_budget_to_locked(void) {
    uint32_t cap = g_stream_expert_cache_entry_count;
    const uint32_t locked_slots = ds4_gpu_stream_expert_slab_locked_slot_count();
    if (locked_slots != 0 && locked_slots < cap) cap = locked_slots;
    const uint64_t gib = 1024ull * 1024ull * 1024ull;
    uint64_t safe_gib = g_stream_expert_cache_mlock_bytes / gib;
    if (safe_gib > 1) safe_gib--;
    if (safe_gib != 0 && g_stream_expert_cache_expert_bytes != 0) {
        uint64_t safe_bytes =
            safe_gib > UINT64_MAX / gib ? UINT64_MAX : safe_gib * gib;
        uint64_t safe_cap64 = safe_bytes / g_stream_expert_cache_expert_bytes;
        if (safe_cap64 > UINT32_MAX) safe_cap64 = UINT32_MAX;
        if (safe_cap64 != 0 && safe_cap64 < cap) {
            cap = (uint32_t)safe_cap64;
        }
    }
    if (cap == 0) return;
    if (g_stream_expert_cache_mlock_budget_cap == 0 ||
        cap < g_stream_expert_cache_mlock_budget_cap) {
        g_stream_expert_cache_mlock_budget_cap = cap;
    }
}

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++;
    if (g_ssd_streaming_mode) {
        void *ptr = [buffer contents];
        const NSUInteger n = [buffer length];
        const double t0 = ds4_gpu_now_ms();
        if (ptr && n != 0 && mlock(ptr, (size_t)n) == 0) {
            const double dt = ds4_gpu_now_ms() - t0;
            g_stream_expert_cache_mlock_ms += dt;
            if (g_stream_expert_cache_mlock_bytes > UINT64_MAX - (uint64_t)n) {
                g_stream_expert_cache_mlock_bytes = UINT64_MAX;
            } else {
                g_stream_expert_cache_mlock_bytes += (uint64_t)n;
            }
        } else {
            const double dt = ds4_gpu_now_ms() - t0;
            g_stream_expert_cache_mlock_ms += dt;
            g_stream_expert_cache_mlock_failures++;
            if (g_stream_expert_cache_mlock_fail_bytes > UINT64_MAX - (uint64_t)n) {
                g_stream_expert_cache_mlock_fail_bytes = UINT64_MAX;
            } else {
                g_stream_expert_cache_mlock_fail_bytes += (uint64_t)n;
            }
            const int err = ptr && n != 0 ? errno : 0;
            ds4_gpu_stream_expert_cache_cap_budget_to_locked();
            ds4_gpu_stream_expert_cache_warn_mlock_failure((uint64_t)n, err);
            if (getenv("DS4_METAL_STREAMING_EXPERT_BUFFER_MLOCK_PROFILE") != NULL) {
                fprintf(stderr,
                        "ds4: Metal streaming expert buffer mlock failed len=%.2f MiB: %s\n",
                        ds4_gpu_mib((uint64_t)n),
                        err != 0 ? strerror(err) : "mlock unavailable");
            }
        }
    }
    return buffer;
}

static int ds4_gpu_stream_expert_combined_buffer_enabled(void) {
    return g_ssd_streaming_mode;
}

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_lock_slot(uint32_t slot) {
    if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
        g_stream_expert_cache_slab_slot_locked[slot]) {
        return 1;
    }
    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] ||
        g_stream_expert_cache_slab_slot_bytes == 0 ||
        base > (uint64_t)NSUIntegerMax) {
        return 0;
    }

    void *contents = [g_stream_expert_cache_slabs[slab] contents];
    if (!contents) return 0;
    const double t0 = ds4_gpu_now_ms();
    void *ptr = (uint8_t *)contents + (NSUInteger)base;
    const size_t n = (size_t)g_stream_expert_cache_slab_slot_bytes;
    if (mlock(ptr, n) == 0) {
        const double dt = ds4_gpu_now_ms() - t0;
        g_stream_expert_cache_mlock_ms += dt;
        if (g_stream_expert_cache_mlock_bytes >
            UINT64_MAX - g_stream_expert_cache_slab_slot_bytes) {
            g_stream_expert_cache_mlock_bytes = UINT64_MAX;
        } else {
            g_stream_expert_cache_mlock_bytes +=
                g_stream_expert_cache_slab_slot_bytes;
        }
        g_stream_expert_cache_slab_slot_locked[slot] = 1;
        return 1;
    }

    const double dt = ds4_gpu_now_ms() - t0;
    g_stream_expert_cache_mlock_ms += dt;
    g_stream_expert_cache_mlock_failures++;
    if (g_stream_expert_cache_mlock_fail_bytes >
        UINT64_MAX - g_stream_expert_cache_slab_slot_bytes) {
        g_stream_expert_cache_mlock_fail_bytes = UINT64_MAX;
    } else {
        g_stream_expert_cache_mlock_fail_bytes +=
            g_stream_expert_cache_slab_slot_bytes;
    }
    ds4_gpu_stream_expert_cache_cap_budget_to_locked();
    ds4_gpu_stream_expert_cache_warn_mlock_failure(
            g_stream_expert_cache_slab_slot_bytes,
            errno);
    return 0;
}

static int ds4_gpu_stream_expert_slab_unlock_slot(uint32_t slot) {
    if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
        !g_stream_expert_cache_slab_slot_locked[slot]) {
        return 1;
    }
    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] ||
        g_stream_expert_cache_slab_slot_bytes == 0 ||
        base > (uint64_t)NSUIntegerMax) {
        return 0;
    }

    void *contents = [g_stream_expert_cache_slabs[slab] contents];
    if (!contents) return 0;
    void *ptr = (uint8_t *)contents + (NSUInteger)base;
    const size_t n = (size_t)g_stream_expert_cache_slab_slot_bytes;
    if (munlock(ptr, n) != 0) return 0;

    g_stream_expert_cache_slab_slot_locked[slot] = 0;
    if (g_stream_expert_cache_mlock_bytes >= g_stream_expert_cache_slab_slot_bytes) {
        g_stream_expert_cache_mlock_bytes -= g_stream_expert_cache_slab_slot_bytes;
    } else {
        g_stream_expert_cache_mlock_bytes = 0;
    }
    return 1;
}

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];
    if (!ds4_gpu_stream_expert_slab_lock_slot(slot)) return 0;
    *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 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 ||
            (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;
    }
    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_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 > 6) {
        return 0;
    }

    const NSUInteger bytes = 6u * 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[6] = {0, 0, 0, 0, 0, 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 (!ds4_gpu_stream_expert_addr_table_requested() &&
        !ds4_gpu_stream_full_expert_addr_table_requested()) {
        return 0;
    }
    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(
        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;
    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_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;
    memset(g_stream_expert_cache_slab_slot_locked,
           0,
           sizeof(g_stream_expert_cache_slab_slot_locked));
    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_mlock_bytes = 0;
        g_stream_expert_cache_mlock_fail_bytes = 0;
        g_stream_expert_cache_mlock_failures = 0;
        g_stream_expert_cache_mlock_ms = 0.0;
        g_stream_expert_cache_mlock_warned = 0;
        g_stream_expert_cache_mlock_budget_cap = 0;
        g_stream_expert_cache_mlock_relief_applied = 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_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_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;
    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 (!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 (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) {
        if (skipped_inflight && !waited_inflight) {
            waited_inflight = 1;
            if (!ds4_gpu_stream_expert_cache_wait_inflight(
                    "streaming expert cache reuse")) {
                return 0;
            }
            goto retry;
        }
        return 0;
    }
    ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer,
                                                     victim_expert,
                                                     1,
                                                     1,
                                                     reuse);
    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_release_mlock_margin(
        uint32_t       protect_layer,
        const int32_t *protect_ids,
        uint32_t       n_protect) {
    if (g_stream_expert_cache_mlock_relief_applied ||
        g_stream_expert_cache_slab_slot_bytes == 0) {
        return 0;
    }

    const uint32_t locked_before =
        ds4_gpu_stream_expert_slab_locked_slot_count();
    if (locked_before == 0) return 0;

    const uint32_t target_release = (locked_before + 9u) / 10u;
    uint32_t released = 0;
    while (released < target_release) {
        uint32_t victim_layer = UINT32_MAX;
        uint32_t victim_expert = UINT32_MAX;
        uint32_t victim_slot = 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 ||
                    !e->slab_backed ||
                    e->slab_slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
                    !g_stream_expert_cache_slab_slot_locked[e->slab_slot] ||
                    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;
                    victim_slot = e->slab_slot;
                }
            }
        }

        if (victim_layer == UINT32_MAX ||
            victim_expert == UINT32_MAX ||
            victim_slot == UINT32_MAX) {
            break;
        }

        ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer,
                                                         victim_expert,
                                                         1,
                                                         0,
                                                         NULL);
        if (ds4_gpu_stream_expert_slab_unlock_slot(victim_slot)) {
            released++;
        }
    }

    if (released == 0) return 0;
    g_stream_expert_cache_mlock_relief_applied = 1;

    uint32_t cap = g_stream_expert_cache_entry_count;
    const uint32_t locked_after =
        ds4_gpu_stream_expert_slab_locked_slot_count();
    if (locked_after != 0 && locked_after < cap) cap = locked_after;
    if (cap != 0 &&
        (g_stream_expert_cache_mlock_budget_cap == 0 ||
         cap < g_stream_expert_cache_mlock_budget_cap)) {
        g_stream_expert_cache_mlock_budget_cap = cap;
    }

    const uint64_t released_bytes =
        (uint64_t)released * g_stream_expert_cache_slab_slot_bytes;
    fprintf(stderr,
            "ds4:   released locked cache margin: %u experts / %.2f GiB; "
            "runtime cache cap now %u experts\n",
            released,
            ds4_gpu_gib(released_bytes),
            g_stream_expert_cache_mlock_budget_cap);
    return released;
}

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_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 (g_stream_expert_cache_mlock_budget_cap != 0) {
            ds4_gpu_stream_expert_cache_release_mlock_margin(protect_layer,
                                                             protect_ids,
                                                             n_protect);
            if (!ds4_gpu_stream_expert_cache_take_reusable(1,
                                                           protect_layer,
                                                           protect_ids,
                                                           n_protect,
                                                           gate_expert_bytes,
                                                           down_expert_bytes,
                                                           &reuse)) {
                return 0;
            }
            *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 *gate_buf && *up_buf && *down_buf;
        }
        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;
        }
        if (g_stream_expert_cache_mlock_budget_cap != 0) {
            ds4_gpu_stream_expert_cache_release_mlock_margin(protect_layer,
                                                             protect_ids,
                                                             n_protect);
            if (!ds4_gpu_stream_expert_cache_take_reusable(1,
                                                           protect_layer,
                                                           protect_ids,
                                                           n_protect,
                                                           gate_expert_bytes,
                                                           down_expert_bytes,
                                                           &reuse)) {
                return 0;
            }
            *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 *gate_buf && *up_buf && *down_buf;
        }
        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_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 < 6; 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[6],
        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);
    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)];
    }

    ds4_gpu_stream_expert_cache_entry *loaded_entries[6] = {
        NULL, NULL, NULL, NULL, NULL, 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 (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[6]) {
    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;
    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[6],
        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 > 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;

    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;
    for (uint32_t i = 0; i < 6; 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;
    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 (!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])) {
            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;
        }
        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,
        };
    }

    const uint32_t n_workers =
        ds4_gpu_stream_expert_pread_thread_count(p->n_tasks);
    p->start_ms = ds4_gpu_now_ms();
    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;
    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[6],
        uint32_t       n_total_expert,
        uint32_t       n_selected,
        const uint64_t gate_abs_offsets[6],
        const uint64_t up_abs_offsets[6],
        const uint64_t down_abs_offsets[6],
        uint64_t       gate_expert_bytes,
        uint64_t       down_expert_bytes,
        uint32_t       missing_mask,
        ds4_gpu_stream_expert_cache_entry *entries[6]) {
    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 > 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 0;
    }
    missing_mask &= (1u << n_selected) - 1u;
    if (missing_mask == 0) return 1;
    if (g_stream_expert_pending_load.active) {
        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();
        }
    }

    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[6] = { 0, 0, 0, 0, 0, 0 };
    uint32_t source_slots[6] = {
        UINT32_MAX, UINT32_MAX, UINT32_MAX,
        UINT32_MAX, UINT32_MAX, 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[6] = { nil, nil, nil, nil, nil, nil };
    __strong id<MTLBuffer> up_bufs[6] = { nil, nil, nil, nil, nil, nil };
    __strong id<MTLBuffer> down_bufs[6] = { nil, nil, nil, nil, nil, nil };
    NSUInteger gate_inners[6] = { 0, 0, 0, 0, 0, 0 };
    NSUInteger up_inners[6] = { 0, 0, 0, 0, 0, 0 };
    NSUInteger down_inners[6] = { 0, 0, 0, 0, 0, 0 };
    ds4_gpu_stream_expert_pread_task tasks[18];
    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;

    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 (!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])) {
            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;

        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) {
        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_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) {
    if (!g_ssd_streaming_mode ||
        !model_map ||
        !selected ||
        !gate_addrs ||
        !up_addrs ||
        !down_addrs ||
        !resources ||
        !n_resources ||
        !unique_out ||
        layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
        n_tokens == 0 ||
        n_selected == 0 ||
        n_selected > 6 ||
        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);
    }
    if (ok) {
        const uint32_t cache_budget =
            ds4_gpu_stream_expert_cache_configured_budget();
        if (cache_budget != 0 && unique_count > cache_budget) {
            fprintf(stderr,
                    "ds4: Metal streaming prefill batch selected addr layer=%u "
                    "needs %u unique experts but cache budget is %u\n",
                    layer,
                    unique_count,
                    cache_budget);
            ok = 0;
        }
    }
    if (ok) {
        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;
            }

            ds4_gpu_stream_expert_cache_entry *entry =
                ds4_gpu_stream_expert_cache_get_protected(model_map,
                                                          model_size,
                                                          layer,
                                                          expert,
                                                          n_total_expert,
                                                          n_selected,
                                                          gate_offset + gate_rel,
                                                          up_offset + gate_rel,
                                                          down_offset + down_rel,
                                                          gate_expert_bytes,
                                                          down_expert_bytes,
                                                          unique_ids,
                                                          unique_count);
            if (!entry) {
                fprintf(stderr,
                        "ds4: Metal streaming prefill batch selected addr failed "
                        "to load layer=%u expert=%u unique=%u budget=%u\n",
                        layer,
                        expert,
                        unique_count,
                        ds4_gpu_stream_expert_cache_configured_budget());
                ok = 0;
                break;
            }
            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) {
        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;
    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) {
    return g_ssd_streaming_mode || 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 = 16u * 32u;
            const NSUInteger k_shared = 32u * 128u;
            const NSUInteger dot_shared = 16u * 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,
        };
        const NSUInteger smem = (((NSUInteger)nth * sizeof(int32_t)) + 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) {
    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 = 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 (n_tok == 1) {
            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;

            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;
        }

        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(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: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);
    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_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) {
    (void)out0; (void)out1; (void)model_map; (void)model_size;
    (void)weight0_offset; (void)weight1_offset;
    (void)in_dim; (void)out0_dim; (void)out1_dim; (void)x; (void)n_tok;
    return 0;
}

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;
}

int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
        ds4_gpu_tensor       *gate,
        ds4_gpu_tensor       *up,
        ds4_gpu_tensor       *mid,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                gate_offset,
        uint64_t                up_offset,
        uint64_t                in_dim,
        uint64_t                out_dim,
        const ds4_gpu_tensor *x,
        float                   clamp) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!gate || !up || !mid || !x || !model_map ||
        (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> 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 = 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 ||
            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 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;
        }

        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;
        id<MTLComputePipelineState> pipeline =
            ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_gate_up_swiglu_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:&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,
                                              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 expert 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_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 (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: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_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 > UINT32_MAX || n_tok != 1) return 0;

    @autoreleasepool {
        id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
        id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
        const uint64_t x_bytes = in_dim * sizeof(float);
        const uint64_t out_bytes = 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_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;

        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;
}

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_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;
}

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;
        }

        uint64_t inner_offset = 0;
        id<MTLBuffer> wbuf = 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_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;
    }

    ds4_gpu_cpy_args f32_to_f16 =
        ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t));
    ds4_gpu_cpy_args f16_to_f32 =
        ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float));
    const NSUInteger nth_f32_f16 = ds4_gpu_cpy_threads(n, g_cpy_f32_f16_pipeline);
    const NSUInteger nth_f16_f32 = ds4_gpu_cpy_threads(n, g_cpy_f16_f32_pipeline);
    const NSUInteger groups_f32_f16 = ((NSUInteger)n + nth_f32_f16 - 1u) / nth_f32_f16;
    const NSUInteger groups_f16_f32 = ((NSUInteger)n + nth_f16_f32 - 1u) / nth_f16_f32;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f32_f16_pipeline];
    [enc setBytes:&f32_to_f16 length:sizeof(f32_to_f16) atIndex:0];
    [enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1];
    [enc setBuffer:g_f16_round_scratch_buffer offset:0 atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(groups_f32_f16, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth_f32_f16, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f16_f32_pipeline];
    [enc setBytes:&f16_to_f32 length:sizeof(f16_to_f32) atIndex:0];
    [enc setBuffer:g_f16_round_scratch_buffer offset:0 atIndex:1];
    [enc setBuffer:g_raw_store_round_buffer offset:0 atIndex:2];
    [enc dispatchThreadgroups:MTLSizeMake(groups_f16_f32, 1, 1)
         threadsPerThreadgroup:MTLSizeMake(nth_f16_f32, 1, 1)];
    ds4_gpu_end_compute_encoder(cb, enc);

    return 1;
}

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;

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

    NSUInteger nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup;
    if (nth_max > 256u) nth_max = 256u;
    NSUInteger nth = 1u;
    while (2u * nth < (NSUInteger)n && nth < nth_max) {
        nth *= 2u;
    }

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_add_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(1, 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 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_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 NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float);
                if (!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_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) {
                    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 && 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) {
                    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);
        if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
                                                    &g_compressor_store_score_bytes,
                                                    score_bytes,
                                                    "ds4_compressor_store_score") ||
                   !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);
        }

        if (ok) {
            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());
        }

        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);

        if (ok) {
            /*
             * 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) {
            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 && 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) {
            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) {
    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;
        }

        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) != 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_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_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) != 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;

    ds4_gpu_cpy_args args =
        ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t));
    const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f32_f16_pipeline);
    const NSUInteger groups = ((NSUInteger)n + 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(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_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;

    ds4_gpu_cpy_args args =
        ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float));
    const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f16_f32_pipeline);
    const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_cpy_f16_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_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 (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_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,
                                              (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;

    id<MTLBuffer> kvbuf = rawbuf;
    NSUInteger kvoff = ds4_gpu_tensor_offset(raw_kv);
    if (raw_start != 0) {
        const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes;
        const uint32_t tail_avail = raw_cap - raw_start;
        const uint32_t tail_rows = tail_avail < n_raw ? tail_avail : n_raw;
        const uint32_t head_rows = n_raw - tail_rows;
        const uint32_t tail_elems = tail_rows * head_dim;
        const uint32_t head_elems = head_rows * head_dim;
        if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer,
                                             &g_flash_attn_ring_bytes,
                                             ring_bytes,
                                             "ds4_flash_attn_ring")) {
            return 0;
        }

        if ((tail_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes,
                                              g_flash_attn_ring_buffer,
                                              0,
                                              tail_elems)) ||
            (head_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              g_flash_attn_ring_buffer,
                                              (NSUInteger)tail_rows * row_bytes,
                                              head_elems))) {
            return 0;
        }

        kvbuf = g_flash_attn_ring_buffer;
        kvoff = 0;
    }

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         kvbuf,
                                         kvoff,
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_raw * head_dim)) {
        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;
}

static void ds4_gpu_fill_raw_prefill_mask(uint16_t *mask, uint32_t n_tokens, uint32_t window) {
    const uint16_t neg_inf_half = 0xfc00u;
    for (uint32_t q = 0; q < n_tokens; q++) {
        uint16_t *row = mask + (uint64_t)q * n_tokens;
        for (uint32_t k = 0; k < n_tokens; k++) {
            const bool causal = k <= q;
            const bool in_window = window == 0 || q - k < window;
            row[k] = causal && in_window ? 0u : neg_inf_half;
        }
    }
}

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;
        }
    }
}

static void ds4_gpu_fill_static_mixed_prefill_mask(
        uint16_t *mask,
        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_tokens; q++) {
        uint16_t *row = mask + (uint64_t)q * n_keys;
        for (uint32_t k = 0; k < n_tokens; k++) {
            const bool causal = k <= q;
            const bool in_window = window == 0 || q - k < window;
            row[k] = causal && in_window ? 0u : neg_inf_half;
        }

        const uint32_t n_visible = (q + 1u) / ratio;
        for (uint32_t c = 0; c < n_comp; c++) {
            row[n_tokens + c] = c < n_visible ? 0u : neg_inf_half;
        }
    }
}

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               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 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_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;
    }

    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_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 &&
        !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");
    }

    ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
                                             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;
    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_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
    }

    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 = nil;
    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_METAL_PROFILE_FLASH_ATTN_STAGE("block_map");

    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);
    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],
                                             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,
                                              (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               n_tokens,
        uint32_t               n_comp,
        uint32_t               window,
        uint32_t               ratio,
        uint32_t               n_head,
        uint32_t               head_dim) {
    if (n_tokens >= 20) {
        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,
                                                                                       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);
}

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               n_tokens,
        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_tokens == 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_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 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_tokens % 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_tokens * (NSUInteger)n_tokens * sizeof(uint16_t);
    const NSUInteger kv_bytes = (NSUInteger)n_tokens * 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_tokens + 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;
    }

    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_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], n_tokens, 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_tokens * 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_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: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_tokens,
        .ne30 = (int32_t)n_tokens,
        .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: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_METAL_PROFILE_FLASH_ATTN_STAGE("block_map");

    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_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 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);
    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               n_tokens,
        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_tokens == 0) {
        return 0;
    }
    if (n_tokens >= 20) {
        return ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(cbp,
                                                                         heads,
                                                                         sinks_buf,
                                                                         sinks_offset,
                                                                         q,
                                                                         raw_kv,
                                                                         n_tokens,
                                                                         window,
                                                                         n_head,
                                                                         head_dim);
    }

    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], 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,
                                              (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);

    if (!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<MTLComputePipelineState> pad_pipeline = nil;
    if ((n_keys % 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_keys % ncpsg) != 0,
                                              (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;

    id<MTLBuffer> raw_linear_buf = rawbuf;
    NSUInteger raw_linear_offset = ds4_gpu_tensor_offset(raw_kv);
    if (raw_start != 0) {
        const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes;
        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;
        if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer,
                                             &g_flash_attn_ring_bytes,
                                             ring_bytes,
                                             "ds4_flash_attn_ring")) {
            return 0;
        }

        if ((tail_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes,
                                              g_flash_attn_ring_buffer,
                                              0,
                                              tail_rows * head_dim)) ||
            (head_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              g_flash_attn_ring_buffer,
                                              (NSUInteger)tail_rows * row_bytes,
                                              head_rows * head_dim))) {
            return 0;
        }

        raw_linear_buf = g_flash_attn_ring_buffer;
        raw_linear_offset = 0;
    }

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         raw_linear_buf,
                                         raw_linear_offset,
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_raw * head_dim)) {
        return 0;
    }
    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_raw * row_bytes_f16,
                                           n_comp * head_dim)) {
            return 0;
        }
    }

    if (!ds4_gpu_encode_fill_f16_1d(cb, g_flash_attn_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),
                                         g_flash_attn_mask_buffer,
                                         (NSUInteger)n_raw * sizeof(uint16_t),
                                         n_comp)) {
        return 0;
    }

    if ((n_keys % ncpsg) != 0) {
        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:g_flash_attn_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:g_flash_attn_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) {
    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;
    }

    id<MTLBuffer> kvbuf = rawbuf;
    NSUInteger kvoff = ds4_gpu_tensor_offset(raw_kv);
    if (raw_start != 0) {
        const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes;
        const uint32_t tail_avail = raw_cap - raw_start;
        const uint32_t tail_rows = tail_avail < n_raw ? tail_avail : n_raw;
        const uint32_t head_rows = n_raw - tail_rows;
        if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer,
                                             &g_flash_attn_ring_bytes,
                                             ring_bytes,
                                             "ds4_flash_attn_ring")) {
            return 0;
        }
        if ((tail_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes,
                                              g_flash_attn_ring_buffer,
                                              0,
                                              tail_rows * head_dim)) ||
            (head_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              g_flash_attn_ring_buffer,
                                              (NSUInteger)tail_rows * row_bytes,
                                              head_rows * head_dim))) {
            return 0;
        }
        kvbuf = g_flash_attn_ring_buffer;
        kvoff = 0;
    }

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         kvbuf,
                                         kvoff,
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_raw * head_dim)) {
        return 0;
    }

    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);
    }
    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;
    }

    id<MTLBuffer> kvbuf = rawbuf;
    NSUInteger kvoff = ds4_gpu_tensor_offset(raw_kv);
    if (raw_start != 0) {
        const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes;
        const uint32_t tail_avail = raw_cap - raw_start;
        const uint32_t tail_rows = tail_avail < n_raw ? tail_avail : n_raw;
        const uint32_t head_rows = n_raw - tail_rows;
        if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer,
                                             &g_flash_attn_ring_bytes,
                                             ring_bytes,
                                             "ds4_flash_attn_ring")) {
            return 0;
        }
        if ((tail_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes,
                                              g_flash_attn_ring_buffer,
                                              0,
                                              tail_rows * head_dim)) ||
            (head_rows &&
             !ds4_gpu_encode_cpy_f32_f32_1d(cb,
                                              rawbuf,
                                              ds4_gpu_tensor_offset(raw_kv),
                                              g_flash_attn_ring_buffer,
                                              (NSUInteger)tail_rows * row_bytes,
                                              head_rows * head_dim))) {
            return 0;
        }
        kvbuf = g_flash_attn_ring_buffer;
        kvoff = 0;
    }

    if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
                                         kvbuf,
                                         kvoff,
                                         g_flash_attn_kv_buffer,
                                         0,
                                         n_raw * head_dim) ||
        !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_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t                n_tokens,
        uint32_t                window,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map || n_tokens == 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,
                                                                n_tokens,
                                                                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_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)) {
            return 0;
        }

        if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode 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_tensor(
        ds4_gpu_tensor       *heads,
        const void             *model_map,
        uint64_t                model_size,
        uint64_t                sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        const ds4_gpu_tensor *comp_kv,
        uint32_t                comp_kv_f16,
        uint32_t                n_tokens,
        uint32_t                n_comp,
        uint32_t                window,
        uint32_t                ratio,
        uint32_t                n_head,
        uint32_t                head_dim) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!heads || !q || !raw_kv || !model_map || 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,
                                                                                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_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,
                                                                                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_pipeline.maxTotalThreadsPerThreadgroup;
        const NSUInteger ds4_nth = n > 1 ? (NSUInteger)n / 2u : 1u;
        if (nth > ds4_nth) nth = ds4_nth;
        if (nth == 0) nth = 1;

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_swiglu_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(1, 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;

        const uint64_t row_bytes = (uint64_t)n * sizeof(float);
        ds4_gpu_bin_args args = {
            .ne00 = (int32_t)n,
            .ne01 = 1,
            .ne02 = 1,
            .ne03 = 1,
            .nb00 = sizeof(float),
            .nb01 = row_bytes,
            .nb02 = row_bytes,
            .nb03 = row_bytes,
            .ne10 = (int32_t)n,
            .ne11 = 1,
            .ne12 = 1,
            .ne13 = 1,
            .nb10 = sizeof(float),
            .nb11 = row_bytes,
            .nb12 = row_bytes,
            .nb13 = row_bytes,
            .ne0 = (int32_t)n,
            .ne1 = 1,
            .ne2 = 1,
            .ne3 = 1,
            .nb0 = sizeof(float),
            .nb1 = row_bytes,
            .nb2 = row_bytes,
            .nb3 = row_bytes,
            .offs = 0,
            .o1 = { 0 },
        };
        NSUInteger nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup;
        if (nth_max > 256u) nth_max = 256u;
        NSUInteger nth = 1;
        while (2u * nth < (NSUInteger)args.ne0 && nth < nth_max) {
            nth *= 2u;
        }

        id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
        [enc setComputePipelineState:g_add_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(1, 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;
}

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_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";
    default:                       return "unknown";
    }
}

static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) {
    if (type == DS4_METAL_TENSOR_IQ2_XXS) {
        return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t);
    }
    return 0;
}

static id<MTLComputePipelineState> ds4_gpu_routed_mv_pipeline(uint32_t type) {
    switch (type) {
    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_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);
    default:
        return nil;
    }
}

static id<MTLComputePipelineState> ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t type) {
    switch (type) {
    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);
    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) {
    if (!cb || !pipeline || !args || !src0 || !src1 || !dst ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
        return 0;
    }

    const NSUInteger rows_per_group = (NSUInteger)args->nr0;
    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,
        NSUInteger                  threadgroup_bytes,
        NSUInteger                  nsg) {
    if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids ||
        args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || 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];
    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) {
    if (!cb || !pipeline || !args || !act || !entries || n_entries == 0 ||
        !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 < 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];
    }
    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) {
    if (!cb || !pipeline || !args || !entries || n_entries == 0 ||
        !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 (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) {
    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:8192u 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);
}

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:4096u 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_pipeline.maxTotalThreadsPerThreadgroup;
    const NSUInteger ds4_nth = n > 1 ? (NSUInteger)n / 2u : 1u;
    if (nth > ds4_nth) nth = ds4_nth;
    if (nth == 0) nth = 1u;

    id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
    [enc setComputePipelineState:g_swiglu_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(1, 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 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;
    }

    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) {
        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(g_dsv4_router_finalize_one_pipeline,
                                    "kernel_dsv4_router_finalize_one");
        id<MTLComputePipelineState> router_weights_pipeline =
            ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline,
                                    "kernel_dsv4_router_weights_one");
        if (!softplus_sqrt_pipeline || !router_finalize_pipeline || !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];
        [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);

        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;
    }

    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;

    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;
    }

    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_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 > 6 || (!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,
        uint32_t                layer_index) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights ||
        n_total_expert == 0 || n_expert == 0 || n_expert > 6 ||
        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;
    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_total_expert * gate_expert_bytes;
        const uint64_t down_tensor_bytes = (uint64_t)n_total_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[6] = { nil, nil, nil, nil, nil, nil };
        __unsafe_unretained id<MTLBuffer> up_slot_bufs[6] = { nil, nil, nil, nil, nil, nil };
        __unsafe_unretained id<MTLBuffer> down_slot_bufs[6] = { nil, nil, nil, nil, nil, nil };
        ds4_gpu_stream_expert_cache_entry *stream_slot_entries[6] = {
            NULL, NULL, NULL, NULL, NULL, NULL
        };
        NSUInteger gate_slot_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        NSUInteger up_slot_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        NSUInteger down_slot_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        uint64_t stream_gate_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        uint64_t stream_up_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
        uint64_t stream_down_abs_offsets[6] = { 0, 0, 0, 0, 0, 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[6] = { 0, 0, 0, 0, 0, 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"))) {
            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);
            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);

        const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type);
        const NSUInteger down_smem = ds4_gpu_routed_mv_smem(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;
        }
        const bool direct_down_sum =
            !g_quality_mode &&
            n_expert == 6 &&
            n_tokens == 1 &&
            down_sum6_pipeline != nil;
        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 =
            !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 =
            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_selected_slots = use_q4_selected_slots || use_iq2_selected_slots;
        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) {
                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) {
                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) 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]) {
                    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")) {
                    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]) {
                    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]) {
                    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_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_q4_selected_slots) &&
                stream_expert_cache_size_known &&
                ds4_gpu_stream_expert_cache_effective_cap(layer_index,
                                                          n_total_expert,
                                                          n_expert) != 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)) {
                    return 0;
                }
                for (uint32_t i = 0; i < 6; 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) {
                    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)) {
                        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) {
                        if (q4_selected_shared_event) {
                            if (ds4_gpu_signal_batch_and_wait_event("selected-id readback") == 0) return 0;
                        } else if (ds4_gpu_end_commands() == 0) {
                            return 0;
                        }
                        if (ds4_gpu_tensor_read(selected,
                                                0,
                                                selected_ids,
                                                (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) {
                            return 0;
                        }
                        if (!q4_selected_shared_event && ds4_gpu_begin_commands() == 0) return 0;
                    } else if (ds4_gpu_tensor_read(selected,
                                                   0,
                                                   selected_ids,
                                                   (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) {
                        return 0;
                    }
                }
            }
            if (selected_timing) {
                selected_read_ms = ds4_gpu_now_ms() - selected_t0;
                selected_t0 = ds4_gpu_now_ms();
            }

            if (selected_ids_available) {
                for (uint32_t i = 0; i < 6; 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);
                        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)) {
                    return 0;
                }

                for (uint32_t i = 0; i < 6; 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]) {
                        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 &&
                    ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
                                                             &stream_gate_addr_buf,
                                                             &stream_up_addr_buf,
                                                             &stream_down_addr_buf);
                use_stream_expert_masked_addr_table =
                    use_stream_expert_addr_table &&
                    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 &&
                    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)) {
                        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)) {
                        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) 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;
                    }
                }
                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)) {
                        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)) {
                        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_read_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_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) return 0;
        }
        if (q4_grouped_boundary || q4_exact_boundary || q4_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_expert_address_table || use_q4_expert_table) &&
            !ds4_gpu_use_model_residency_set(cb)) {
            return 0;
        }
        if (!q4_table_queue_residency &&
            q4_table_layer_residency &&
            [cb respondsToSelector:@selector(useResidencySet:)]) {
            [cb useResidencySet:q4_table_layer_residency];
        }

        const char *moe_one_stage_layer_env = getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE_LAYER");
        bool moe_one_stage_layer_match = true;
        if (moe_one_stage_layer_env && moe_one_stage_layer_env[0]) {
            char *end = NULL;
            long layer = strtol(moe_one_stage_layer_env, &end, 10);
            moe_one_stage_layer_match =
                end && *end == '\0' && layer >= 0 && (uint32_t)layer == layer_index;
        }
        const bool moe_one_stage_profile =
            getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") != NULL &&
            g_batch_cb != nil &&
            moe_one_stage_layer_match;
        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) {
                return 0;
            }
            cb = ds4_gpu_command_buffer(&owned);
            if (!cb) 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) 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);
                }
            } 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,
                                            2,
                                            false) &&
                 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,
                                            2,
                                            false);
        }
        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 (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);
                }
            } 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),
                                                 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,
                                                 2,
                                                 false);
        }
        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) return 0;

        if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) return 0;
        if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) {
            if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
                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) {
    if (!g_initialized && !ds4_gpu_init()) return 0;
    if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights ||
        n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || n_expert > 6) {
        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_total_expert * gate_expert_bytes;
    const uint64_t down_tensor_bytes = (uint64_t)n_total_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,
                                             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> 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;
        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);
        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);
        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 (usually 2 tokens) 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 <= 4u &&
            !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));
        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;
        const bool use_mm_id_pair_swiglu =
            use_mm_id &&
            request_mid_f16 &&
            gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
            down_type == DS4_METAL_TENSOR_Q2_K &&
            n_expert == 6 &&
            getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL &&
            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));

            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("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)) {
                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;
        }

        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 char *moe_stage_layer_env =
            getenv("DS4_METAL_MOE_STAGE_PROFILE_LAYER");
        uint32_t moe_stage_layer = UINT32_MAX;
        if (moe_stage_layer_env && moe_stage_layer_env[0]) {
            char *end = NULL;
            unsigned long v = strtoul(moe_stage_layer_env, &end, 10);
            if (end != moe_stage_layer_env && *end == '\0' && v <= UINT32_MAX) {
                moe_stage_layer = (uint32_t)v;
            }
        }
        const bool moe_stage_profile =
            getenv("DS4_METAL_MOE_STAGE_PROFILE") != NULL &&
            g_batch_cb != nil &&
            (moe_stage_layer == UINT32_MAX || moe_stage_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_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);
        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);
            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));
                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));
                DS4_METAL_PROFILE_MOE_STAGE("up");
            }
        } 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,
                                                  2,
                                                  false) &&
                 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,
                                                  2,
                                                  false);
        }
        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_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);
                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),
                                                     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);
            } 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,
                                                     2,
                                                     false);
            }
        }
        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) 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_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;
}

/* Decode-only HC-pre plus the immediately following weighted RMSNorm.  This is
 * specialized for DS4's HC=4 decode shape; larger batched prefill keeps using
 * the existing two-stage path. */
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 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;
}

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;
}