armorpaint/base/tools/iris.c/iris_qwen3.c

/*
 * Qwen3 Text Encoder Implementation
 *
 * Implements Qwen3-4B model for text encoding in Iris image generation.
 * - 36 transformer layers
 * - 2560 hidden dimension
 * - GQA with 32 query heads and 8 KV heads
 * - RoPE positional embeddings
 * - SwiGLU MLP
 */

#include "iris_qwen3.h"
#include "iris_gguf.h"
#include "iris_kernels.h"
#include "iris_safetensors.h"
#include <dirent.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* Use Metal for GPU acceleration */
#ifdef USE_METAL
#include "iris_metal.h"
#endif

/* Use Vulkan compute for GPU acceleration (GEMM offload) */
#ifdef USE_VULKAN
#include "iris_vulkan.h"
/* Vulkan dispatch overhead is low, so offload smaller matrices than Metal:
 * the per-layer projections (seq*out ~ 1.3M) are the text-encoder bottleneck. */
#define QWEN3_VK_MIN_ELEMENTS (256 * 1024)
#endif

/* Minimum matrix size for GPU acceleration.
 * Using 10M threshold keeps text encoder on CPU (Accelerate BLAS), which is
 * faster and avoids GPU memory pressure on 16GB systems. Text encoder weights
 * are only used once per generation, so GPU caching provides no benefit.
 * Fixes issue #9: SIGKILL on 16GB Metal systems during text encoding. */
#define QWEN3_MIN_GPU_ELEMENTS (10 * 1024 * 1024)

/* Maximum number of safetensors shards (Qwen3-8B may have >2) */
#define QWEN3_MAX_SHARDS 16

/* ========================================================================
 * Data Structures
 * ======================================================================== */

typedef struct {
	float *q_proj_weight; /* [num_heads * head_dim, hidden] = [4096, 2560] */
	float *k_proj_weight; /* [num_kv_heads * head_dim, hidden] = [1024, 2560] */
	float *v_proj_weight; /* [num_kv_heads * head_dim, hidden] = [1024, 2560] */
	float *o_proj_weight; /* [hidden, num_heads * head_dim] = [2560, 4096] */
	float *q_norm_weight; /* [head_dim] = [128] */
	float *k_norm_weight; /* [head_dim] = [128] */
	/* BF16 weight pointers (for GPU path) */
	uint16_t *q_proj_weight_bf16;
	uint16_t *k_proj_weight_bf16;
	uint16_t *v_proj_weight_bf16;
	uint16_t *o_proj_weight_bf16;
	uint16_t *q_norm_weight_bf16; /* [head_dim] = [128] */
	uint16_t *k_norm_weight_bf16; /* [head_dim] = [128] */
} qwen3_attention_t;

typedef struct {
	float *gate_proj_weight; /* [intermediate, hidden] = [9728, 2560] */
	float *up_proj_weight;   /* [intermediate, hidden] = [9728, 2560] */
	float *down_proj_weight; /* [hidden, intermediate] = [2560, 9728] */
	/* BF16 weight pointers (for GPU path) */
	uint16_t *gate_proj_weight_bf16;
	uint16_t *up_proj_weight_bf16;
	uint16_t *down_proj_weight_bf16;
} qwen3_mlp_t;

typedef struct {
	float            *input_layernorm_weight;          /* [hidden] */
	float            *post_attention_layernorm_weight; /* [hidden] */
	qwen3_attention_t attn;
	qwen3_mlp_t       mlp;
	/* BF16 layer norm weights (for GPU path) - unused currently, kept for future */
	uint16_t *input_layernorm_weight_bf16;
	uint16_t *post_attention_layernorm_weight_bf16;
} qwen3_layer_t;

struct qwen3_model {
	/* Architecture (from config.json) */
	int   hidden_size;
	int   intermediate_size;
	int   num_heads;
	int   num_kv_heads;
	int   head_dim;
	int   vocab_size;
	float rope_theta;
	int   text_dim; /* 3 * hidden_size (layers 8,17,26 concatenated) */

	/* Embedding layer */
	float *embed_tokens; /* [vocab_size, hidden] */

	/* Transformer layers */
	qwen3_layer_t *layers; /* [num_layers] */
	int            num_layers;

	/* Final layer norm */
	float *norm_weight; /* [hidden] */

	/* RoPE precomputed */
	float *rope_cos; /* [max_seq_len, head_dim/2] */
	float *rope_sin; /* [max_seq_len, head_dim/2] */

	/* Working memory */
	float *hidden_state; /* [seq_len, hidden] */
	float *residual;     /* [seq_len, hidden] */
	float *q_buf;        /* [seq_len, num_heads * head_dim] */
	float *k_buf;        /* [seq_len, num_kv_heads * head_dim] */
	float *v_buf;        /* [seq_len, num_kv_heads * head_dim] */
	float *attn_scores;  /* [num_heads, seq_len, seq_len] */
	float *attn_out;     /* [seq_len, num_heads * head_dim] */
	float *mlp_gate;     /* [seq_len, intermediate] */
	float *mlp_up;       /* [seq_len, intermediate] */
	float *mlp_out;      /* [seq_len, hidden] */
	float *norm_buf;     /* [seq_len, hidden] */

	/* Output layers storage (for extracting layers 9, 18, 27) */
	float *layer_outputs[3]; /* [seq_len, hidden] each */

	/* Pre-allocated attention work buffers (avoid per-call allocation) */
	float *attn_q_head;   /* [seq_len, head_dim] */
	float *attn_v_head;   /* [seq_len, head_dim] */
	float *attn_out_head; /* [seq_len, head_dim] */

	/* Mmap mode: keep safetensors files open, load layer weights on-demand */
	int                 use_mmap;
	safetensors_file_t *sf_files[QWEN3_MAX_SHARDS];
	int                 num_sf_files;

	/* BF16 GPU acceleration */
	int use_bf16;

	/* Set when the large projection weights are GGML Q8_0 block-quantized
	 * (rather than bf16). The resident Vulkan linear ops then dequantize the
	 * Q8_0 weights in-shader. Norms/embeddings remain f32/bf16. */
	int weights_q8;
};

/* Forward declarations for mmap streaming mode */
static int load_layer_weights(qwen3_layer_t *layer, safetensors_file_t **files, int num_files, int layer_idx);
#ifdef USE_METAL
static int load_layer_weights_small_f32(qwen3_layer_t *layer, safetensors_file_t **files, int num_files, int layer_idx);
#endif
#if defined(USE_METAL) || defined(USE_VULKAN)
static int load_layer_weights_bf16(qwen3_layer_t *layer, safetensors_file_t **files, int num_files, int layer_idx);
#endif
static void free_layer_weights(qwen3_layer_t *layer);

/* ========================================================================
 * Basic Operations
 * ======================================================================== */

static void qwen3_linear(float *y, const float *x, const float *W, int seq_len, int in_dim, int out_dim) {
	/* y[seq, out] = x[seq, in] @ W[out, in]^T */
#ifdef USE_METAL
	/* Use GPU for large matrices */
	size_t matrix_elements = (size_t)seq_len * out_dim;
	if (iris_metal_available() && matrix_elements >= QWEN3_MIN_GPU_ELEMENTS) {
		iris_metal_sgemm_cached(0, 1, /* no transpose A, transpose B */
		                        seq_len, out_dim, in_dim, 1.0f, x, in_dim, W, in_dim, 0.0f, y, out_dim);
		return;
	}
#endif

#ifdef USE_VULKAN
	if (iris_vulkan_available() && (size_t)seq_len * out_dim >= QWEN3_VK_MIN_ELEMENTS) {
		/* Text-encoder weights are used once per generation, so don't cache
		 * them on the GPU (would needlessly accumulate gigabytes of VRAM).
		 * The non-cached path streams each weight through scratch VRAM. */
		iris_vulkan_sgemm(0, 1, /* no transpose A, transpose B */
		                  seq_len, out_dim, in_dim, 1.0f, x, in_dim, W, in_dim, 0.0f, y, out_dim);
	}
	else
#endif
		for (int s = 0; s < seq_len; s++) {
			for (int o = 0; o < out_dim; o++) {
				float sum = 0.0f;
				for (int i = 0; i < in_dim; i++) {
					sum += x[s * in_dim + i] * W[o * in_dim + i];
				}
				y[s * out_dim + o] = sum;
			}
		}
}

static void qwen3_rms_norm(float *out, const float *x, const float *weight, int seq_len, int hidden, float eps) {
	for (int s = 0; s < seq_len; s++) {
		const float *x_row   = x + s * hidden;
		float       *out_row = out + s * hidden;

		/* Compute RMS */
		float sum_sq = 0.0f;
		for (int i = 0; i < hidden; i++) {
			sum_sq += x_row[i] * x_row[i];
		}
		float rms     = sqrtf(sum_sq / hidden + eps);
		float rms_inv = 1.0f / rms;

		/* Normalize and scale */
		for (int i = 0; i < hidden; i++) {
			out_row[i] = x_row[i] * rms_inv * weight[i];
		}
	}
}

/* Per-head RMS norm for Q/K normalization */
static void qwen3_head_rms_norm(float *out, const float *x, const float *weight, int seq_len, int num_heads, int head_dim, float eps) {
	for (int s = 0; s < seq_len; s++) {
		for (int h = 0; h < num_heads; h++) {
			const float *x_head   = x + s * num_heads * head_dim + h * head_dim;
			float       *out_head = out + s * num_heads * head_dim + h * head_dim;

			/* Compute RMS for this head */
			float sum_sq = 0.0f;
			for (int i = 0; i < head_dim; i++) {
				sum_sq += x_head[i] * x_head[i];
			}
			float rms     = sqrtf(sum_sq / head_dim + eps);
			float rms_inv = 1.0f / rms;

			/* Normalize and scale */
			for (int i = 0; i < head_dim; i++) {
				out_head[i] = x_head[i] * rms_inv * weight[i];
			}
		}
	}
}

static void qwen3_softmax(float *x, int len) {
	float max_val = x[0];
	for (int i = 1; i < len; i++) {
		if (x[i] > max_val)
			max_val = x[i];
	}

	float sum = 0.0f;
	for (int i = 0; i < len; i++) {
		x[i] = fast_expf(x[i] - max_val);
		sum += x[i];
	}

	float inv_sum = 1.0f / sum;
	for (int i = 0; i < len; i++) {
		x[i] *= inv_sum;
	}
}

/* ========================================================================
 * RoPE (Rotary Position Embedding)
 * ======================================================================== */

static void compute_rope_freqs(float *cos_out, float *sin_out, int max_seq_len, int head_dim, float theta) {
	int half_dim = head_dim / 2;

	for (int pos = 0; pos < max_seq_len; pos++) {
		for (int i = 0; i < half_dim; i++) {
			float freq                  = 1.0f / powf(theta, (float)(2 * i) / head_dim);
			float angle                 = pos * freq;
			cos_out[pos * half_dim + i] = cosf(angle);
			sin_out[pos * half_dim + i] = sinf(angle);
		}
	}
}

/* Apply RoPE rotation to Q and K for all attention heads. Uses split-half
 * rotation (dims [0..63] paired with [64..127]) within GQA: each of the
 * 32 Q heads is rotated independently, while only 8 K heads are rotated
 * (shared across groups of 4 Q heads). Position encoding creates
 * relative-position-dependent attention logit decay. */
static void apply_rope(float *q, float *k, const float *cos_cache, const float *sin_cache, int seq_len, int num_q_heads, int num_kv_heads, int head_dim) {
	int half_dim = head_dim / 2;

	/* Apply RoPE to Q */
	for (int s = 0; s < seq_len; s++) {
		const float *cos_row = cos_cache + s * half_dim;
		const float *sin_row = sin_cache + s * half_dim;

		for (int h = 0; h < num_q_heads; h++) {
			float *q_head = q + s * num_q_heads * head_dim + h * head_dim;

			for (int i = 0; i < half_dim; i++) {
				float x0      = q_head[i];
				float x1      = q_head[i + half_dim];
				float cos_val = cos_row[i];
				float sin_val = sin_row[i];

				q_head[i]            = x0 * cos_val - x1 * sin_val;
				q_head[i + half_dim] = x0 * sin_val + x1 * cos_val;
			}
		}
	}

	/* Apply RoPE to K */
	for (int s = 0; s < seq_len; s++) {
		const float *cos_row = cos_cache + s * half_dim;
		const float *sin_row = sin_cache + s * half_dim;

		for (int h = 0; h < num_kv_heads; h++) {
			float *k_head = k + s * num_kv_heads * head_dim + h * head_dim;

			for (int i = 0; i < half_dim; i++) {
				float x0      = k_head[i];
				float x1      = k_head[i + half_dim];
				float cos_val = cos_row[i];
				float sin_val = sin_row[i];

				k_head[i]            = x0 * cos_val - x1 * sin_val;
				k_head[i + half_dim] = x0 * sin_val + x1 * cos_val;
			}
		}
	}
}

/* ========================================================================
 * Attention
 * ======================================================================== */

/* Full attention forward pass for one Qwen3 layer. Implements GQA (Grouped
 * Query Attention): 32 Q heads share 8 KV heads (4:1 ratio for 4B model).
 * Pipeline: Q/K/V projections -> per-head RMSNorm on Q and K -> RoPE ->
 * scaled dot-product attention with causal + padding mask -> output projection.
 * Each Q head attends to its corresponding KV group (h / heads_per_kv). */
static void qwen3_attention_forward(qwen3_model_t *model, qwen3_layer_t *layer, int seq_len, const int *attention_mask) {
	int   num_heads    = model->num_heads;
	int   num_kv_heads = model->num_kv_heads;
	int   head_dim     = model->head_dim;
	int   hidden       = model->hidden_size;
	int   kv_dim       = num_kv_heads * head_dim;
	int   q_dim        = num_heads * head_dim;
	float scale        = 1.0f / sqrtf((float)head_dim);

	/* Q, K, V projections */
	qwen3_linear(model->q_buf, model->norm_buf, layer->attn.q_proj_weight, seq_len, hidden, q_dim);
	qwen3_linear(model->k_buf, model->norm_buf, layer->attn.k_proj_weight, seq_len, hidden, kv_dim);
	qwen3_linear(model->v_buf, model->norm_buf, layer->attn.v_proj_weight, seq_len, hidden, kv_dim);

	/* Q/K RMS normalization (per-head) */
	qwen3_head_rms_norm(model->q_buf, model->q_buf, layer->attn.q_norm_weight, seq_len, num_heads, head_dim, QWEN3_RMS_NORM_EPS);
	qwen3_head_rms_norm(model->k_buf, model->k_buf, layer->attn.k_norm_weight, seq_len, num_kv_heads, head_dim, QWEN3_RMS_NORM_EPS);

	/* Apply RoPE */
	apply_rope(model->q_buf, model->k_buf, model->rope_cos, model->rope_sin, seq_len, num_heads, num_kv_heads, head_dim);

#ifdef USE_METAL
	/* Try GPU-accelerated causal attention for all heads in parallel.
	 * The GPU kernel uses both causal masking and attention mask.
	 * This ensures exact parity with CPU implementation. */
	if (iris_metal_available()) {
		if (iris_metal_causal_attention(model->attn_out, model->q_buf, model->k_buf, model->v_buf, attention_mask, seq_len, num_heads, num_kv_heads, head_dim,
		                                scale)) {
			/* GPU attention succeeded - skip to output projection */
			goto output_proj;
		}
	}
#endif

	/* CPU fallback: compute attention for each head with GQA
	 * Use BLAS for Q@K^T and scores@V matrix multiplications */
	{
		int heads_per_kv = num_heads / num_kv_heads;

		for (int h = 0; h < num_heads; h++) {
			int    kv_h   = h / heads_per_kv; /* Which KV head to use */
			float *scores = model->attn_scores + h * seq_len * seq_len;

			/* Q accessed directly with strided lda (avoids copy)
			 * Q[s,d] = q_buf[s * q_dim + h * head_dim + d] */
			const float *q_strided = model->q_buf + h * head_dim;

			/* K accessed directly with strided lda + CblasTrans (avoids transpose)
			 * K[s,d] = k_buf[s * kv_dim + kv_h * head_dim + d] */
			const float *k_strided = model->k_buf + kv_h * head_dim;

			/* scores = scale * Q @ K^T (strided views, contiguous output) */
#ifdef USE_VULKAN
			if (iris_vulkan_available() && (size_t)seq_len * seq_len >= QWEN3_VK_MIN_ELEMENTS) {
				iris_vulkan_sgemm(0, 1, seq_len, seq_len, head_dim, scale, q_strided, q_dim, k_strided, kv_dim, 0.0f, scores, seq_len);
			}
			else
#endif
				/* Fallback: naive matmul */
				for (int i = 0; i < seq_len; i++) {
					for (int j = 0; j < seq_len; j++) {
						float dot = 0.0f;
						for (int d = 0; d < head_dim; d++) {
							dot += q_strided[i * q_dim + d] * k_strided[j * kv_dim + d];
						}
						scores[i * seq_len + j] = dot * scale;
					}
				}

			/* Apply causal mask and attention mask, then softmax */
			for (int i = 0; i < seq_len; i++) {
				for (int j = 0; j < seq_len; j++) {
					if (j > i) {
						scores[i * seq_len + j] = -1e9f;
					}
					if (attention_mask && attention_mask[j] == 0) {
						scores[i * seq_len + j] = -1e9f;
					}
				}
				qwen3_softmax(scores + i * seq_len, seq_len);
			}

			/* V can be accessed directly with strided lda (avoids copy)
			 * V[s,d] = v_buf[s * kv_dim + kv_h * head_dim + d] */
			const float *v_strided = model->v_buf + kv_h * head_dim;

			/* Output can be written directly with strided ldc (avoids copy)
			 * out[s,d] = attn_out[s * q_dim + h * head_dim + d] */
			float *out_strided = model->attn_out + h * head_dim;

			/* out = scores @ V (strided V input and strided output view)
			 * scores: [seq_len, seq_len], V: [seq_len, head_dim] with ldb=kv_dim */
#ifdef USE_VULKAN
			if (iris_vulkan_available() && (size_t)seq_len * seq_len >= QWEN3_VK_MIN_ELEMENTS) {
				iris_vulkan_sgemm(0, 0, seq_len, head_dim, seq_len, 1.0f, scores, seq_len, v_strided, kv_dim, 0.0f, out_strided, q_dim);
			}
			else
#endif
				for (int i = 0; i < seq_len; i++) {
					for (int d = 0; d < head_dim; d++) {
						float sum = 0.0f;
						for (int j = 0; j < seq_len; j++) {
							sum += scores[i * seq_len + j] * v_strided[j * kv_dim + d];
						}
						out_strided[i * q_dim + d] = sum;
					}
				}
		}
	}

	/* Work buffers are pre-allocated in model, no free needed */

#ifdef USE_METAL
output_proj:
#endif
	/* Output projection */
	qwen3_linear(model->hidden_state, model->attn_out, layer->attn.o_proj_weight, seq_len, q_dim, hidden);
}

/* ========================================================================
 * MLP (SwiGLU)
 * ======================================================================== */

/* SwiGLU MLP: gate = W_gate @ x, up = W_up @ x, out = W_down @ (silu(gate) * up).
 * The gated activation lets the network learn which features to pass through:
 * silu(gate) acts as a learned soft switch on the up-projected features. */
static void qwen3_mlp_forward(qwen3_model_t *model, qwen3_layer_t *layer, int seq_len) {
	int hidden       = model->hidden_size;
	int intermediate = model->intermediate_size;

	/* Gate and Up projections */
	qwen3_linear(model->mlp_gate, model->norm_buf, layer->mlp.gate_proj_weight, seq_len, hidden, intermediate);
	qwen3_linear(model->mlp_up, model->norm_buf, layer->mlp.up_proj_weight, seq_len, hidden, intermediate);

	/* SwiGLU: silu(gate) * up - fused for better performance */
	iris_silu_mul(model->mlp_gate, model->mlp_up, seq_len * intermediate);

	/* Down projection */
	qwen3_linear(model->mlp_out, model->mlp_gate, layer->mlp.down_proj_weight, seq_len, intermediate, hidden);
}

/* ========================================================================
 * Transformer Layer
 * ======================================================================== */

/* One Qwen3 transformer layer (pre-norm architecture):
 * RMSNorm -> self-attention -> residual add -> RMSNorm -> SwiGLU MLP ->
 * residual add. Pre-norm (normalizing before each sub-layer rather than
 * after) improves training stability and is standard in modern LLMs. */
static void qwen3_layer_forward(qwen3_model_t *model, qwen3_layer_t *layer, int seq_len, const int *attention_mask) {
	int hidden = model->hidden_size;

	/* Save residual */
	memcpy(model->residual, model->hidden_state, seq_len * hidden * sizeof(float));

	/* Pre-attention LayerNorm */
	qwen3_rms_norm(model->norm_buf, model->hidden_state, layer->input_layernorm_weight, seq_len, hidden, QWEN3_RMS_NORM_EPS);

	/* Self-attention */
	qwen3_attention_forward(model, layer, seq_len, attention_mask);

	/* Residual connection */
	for (int i = 0; i < seq_len * hidden; i++) {
		model->hidden_state[i] += model->residual[i];
	}

	/* Save residual */
	memcpy(model->residual, model->hidden_state, seq_len * hidden * sizeof(float));

	/* Pre-MLP LayerNorm */
	qwen3_rms_norm(model->norm_buf, model->hidden_state, layer->post_attention_layernorm_weight, seq_len, hidden, QWEN3_RMS_NORM_EPS);

	/* MLP */
	qwen3_mlp_forward(model, layer, seq_len);

	/* Residual connection */
	for (int i = 0; i < seq_len * hidden; i++) {
		model->hidden_state[i] = model->residual[i] + model->mlp_out[i];
	}
}

#ifdef USE_METAL
/* ========================================================================
 * BF16 GPU-Accelerated Layer Forward
 * Uses GPU for linear layers, keeps attention/norm on CPU for simplicity.
 * ======================================================================== */

/* Helper to convert f32 array to bf16 GPU tensor */
static iris_gpu_tensor_t f32_to_bf16_tensor(const float *data, int n) {
	iris_gpu_tensor_t f32_tensor = iris_gpu_tensor_create(data, n);
	if (!f32_tensor)
		return NULL;
	iris_gpu_tensor_t bf16_tensor = iris_gpu_tensor_f32_to_bf16(f32_tensor);
	iris_gpu_tensor_free(f32_tensor);
	return bf16_tensor;
}

/* Helper to read bf16 GPU tensor back to f32 array */
static void bf16_tensor_to_f32(iris_gpu_tensor_t bf16_tensor, float *out) {
	iris_gpu_tensor_t f32_tensor = iris_gpu_tensor_bf16_to_f32(bf16_tensor);
	if (f32_tensor) {
		iris_gpu_tensor_read(f32_tensor, out);
		iris_gpu_tensor_free(f32_tensor);
	}
}

/* Convert bf16 value to f32 */
static inline float bf16_to_f32_val(uint16_t bf16) {
	uint32_t f32_bits = ((uint32_t)bf16) << 16;
	float    result;
	memcpy(&result, &f32_bits, sizeof(float));
	return result;
}

/* Helper to create bf16 GPU tensor from bf16 CPU data (for small tensors like norm weights) */
static iris_gpu_tensor_t bf16_ptr_to_bf16_tensor(const uint16_t *bf16_data, int n) {
	/* Convert bf16->f32 on CPU, then f32->bf16 on GPU */
	float *f32_tmp = malloc(n * sizeof(float));
	if (!f32_tmp)
		return NULL;
	for (int i = 0; i < n; i++) {
		f32_tmp[i] = bf16_to_f32_val(bf16_data[i]);
	}
	iris_gpu_tensor_t result = f32_to_bf16_tensor(f32_tmp, n);
	free(f32_tmp);
	return result;
}

/* GPU-accelerated MLP using bf16 weights */
static void qwen3_mlp_forward_bf16(qwen3_model_t *model, qwen3_layer_t *layer, int seq_len) {
	int hidden       = model->hidden_size;
	int intermediate = model->intermediate_size;
	int n            = seq_len * intermediate;

	/* Convert input to bf16 tensor on GPU */
	iris_gpu_tensor_t x = f32_to_bf16_tensor(model->norm_buf, seq_len * hidden);
	if (!x) {
		qwen3_mlp_forward(model, layer, seq_len);
		return;
	}

	/* Gate and Up projections on GPU */
	iris_gpu_tensor_t gate = iris_gpu_linear_bf16_native(x, layer->mlp.gate_proj_weight_bf16, seq_len, hidden, intermediate);
	iris_gpu_tensor_t up   = iris_gpu_linear_bf16_native(x, layer->mlp.up_proj_weight_bf16, seq_len, hidden, intermediate);
	iris_gpu_tensor_free(x);

	if (!gate || !up) {
		if (gate)
			iris_gpu_tensor_free(gate);
		if (up)
			iris_gpu_tensor_free(up);
		qwen3_mlp_forward(model, layer, seq_len);
		return;
	}

	/* SwiGLU: silu(gate) * up on GPU */
	iris_gpu_silu_mul_bf16(gate, up, n);
	iris_gpu_tensor_free(up);

	/* Down projection on GPU */
	iris_gpu_tensor_t out = iris_gpu_linear_bf16_native(gate, layer->mlp.down_proj_weight_bf16, seq_len, intermediate, hidden);
	iris_gpu_tensor_free(gate);

	if (!out) {
		qwen3_mlp_forward(model, layer, seq_len);
		return;
	}

	/* Read result back to CPU */
	bf16_tensor_to_f32(out, model->mlp_out);
	iris_gpu_tensor_free(out);
}

/* GPU-accelerated attention using bf16 weights for projections */
static void qwen3_attention_forward_bf16(qwen3_model_t *model, qwen3_layer_t *layer, int seq_len, const int *attention_mask) {
	int   num_heads    = model->num_heads;
	int   num_kv_heads = model->num_kv_heads;
	int   head_dim     = model->head_dim;
	int   hidden       = model->hidden_size;
	int   kv_dim       = num_kv_heads * head_dim;
	int   q_dim        = num_heads * head_dim;
	float scale        = 1.0f / sqrtf((float)head_dim);

	/* Convert input to bf16 tensor */
	iris_gpu_tensor_t x = f32_to_bf16_tensor(model->norm_buf, seq_len * hidden);
	if (!x) {
		qwen3_attention_forward(model, layer, seq_len, attention_mask);
		return;
	}

	/* Q, K, V projections on GPU */
	iris_gpu_tensor_t q = iris_gpu_linear_bf16_native(x, layer->attn.q_proj_weight_bf16, seq_len, hidden, q_dim);
	iris_gpu_tensor_t k = iris_gpu_linear_bf16_native(x, layer->attn.k_proj_weight_bf16, seq_len, hidden, kv_dim);
	iris_gpu_tensor_t v = iris_gpu_linear_bf16_native(x, layer->attn.v_proj_weight_bf16, seq_len, hidden, kv_dim);
	iris_gpu_tensor_free(x);

	if (!q || !k || !v) {
		if (q)
			iris_gpu_tensor_free(q);
		if (k)
			iris_gpu_tensor_free(k);
		if (v)
			iris_gpu_tensor_free(v);
		qwen3_attention_forward(model, layer, seq_len, attention_mask);
		return;
	}

	/* Try full bf16 pipeline: Q/K norm, RoPE, and attention all on GPU */
	iris_gpu_tensor_t attn_out = iris_gpu_tensor_alloc_f16(seq_len * q_dim);
	if (attn_out && layer->attn.q_norm_weight_bf16 && layer->attn.k_norm_weight_bf16) {
		/* Get bf16 weight tensors for Q/K norm */
		iris_gpu_tensor_t q_norm_w = bf16_ptr_to_bf16_tensor(layer->attn.q_norm_weight_bf16, head_dim);
		iris_gpu_tensor_t k_norm_w = bf16_ptr_to_bf16_tensor(layer->attn.k_norm_weight_bf16, head_dim);

		if (q_norm_w && k_norm_w) {

			/* Q/K RMS normalization on GPU - separate calls for GQA (different head counts) */
			int q_norm_ok = iris_gpu_head_rms_norm_bf16(q, q_norm_w, seq_len, num_heads, head_dim, QWEN3_RMS_NORM_EPS);
			int k_norm_ok = iris_gpu_head_rms_norm_bf16(k, k_norm_w, seq_len, num_kv_heads, head_dim, QWEN3_RMS_NORM_EPS);

			iris_gpu_tensor_free(q_norm_w);
			iris_gpu_tensor_free(k_norm_w);

			if (q_norm_ok && k_norm_ok) {
				/* Apply RoPE - GPU bf16 */
				iris_gpu_rope_text_bf16(q, k, model->rope_cos, model->rope_sin, seq_len, num_heads, num_kv_heads, head_dim);

				/* GPU causal attention (bf16) */
				if (iris_gpu_causal_attention_bf16(attn_out, q, k, v, attention_mask, seq_len, num_heads, num_kv_heads, head_dim, scale)) {
					iris_gpu_tensor_free(q);
					iris_gpu_tensor_free(k);
					iris_gpu_tensor_free(v);

					/* Output projection on GPU - input already bf16 */
					iris_gpu_tensor_t out = iris_gpu_linear_bf16_native(attn_out, layer->attn.o_proj_weight_bf16, seq_len, q_dim, hidden);
					iris_gpu_tensor_free(attn_out);

					if (out) {
						bf16_tensor_to_f32(out, model->hidden_state);
						iris_gpu_tensor_free(out);
						return; /* Success - full bf16 pipeline */
					}
				}
			}
		}
		else {
			/* q_norm_w or k_norm_w allocation failed */
			if (q_norm_w)
				iris_gpu_tensor_free(q_norm_w);
			if (k_norm_w)
				iris_gpu_tensor_free(k_norm_w);
		}
		/* GPU bf16 path failed - free everything and use full CPU fallback */
		iris_gpu_tensor_free(q);
		iris_gpu_tensor_free(k);
		iris_gpu_tensor_free(v);
		iris_gpu_tensor_free(attn_out);
		qwen3_attention_forward(model, layer, seq_len, attention_mask);
		return;
	}

	/* Fallback: No attn_out or no bf16 norm weights - use CPU path */
	if (attn_out)
		iris_gpu_tensor_free(attn_out);
	bf16_tensor_to_f32(q, model->q_buf);
	bf16_tensor_to_f32(k, model->k_buf);
	bf16_tensor_to_f32(v, model->v_buf);
	iris_gpu_tensor_free(q);
	iris_gpu_tensor_free(k);
	iris_gpu_tensor_free(v);

	/* Q/K RMS normalization (per-head) - CPU */
	qwen3_head_rms_norm(model->q_buf, model->q_buf, layer->attn.q_norm_weight, seq_len, num_heads, head_dim, QWEN3_RMS_NORM_EPS);
	qwen3_head_rms_norm(model->k_buf, model->k_buf, layer->attn.k_norm_weight, seq_len, num_kv_heads, head_dim, QWEN3_RMS_NORM_EPS);

	/* Apply RoPE - CPU */
	apply_rope(model->q_buf, model->k_buf, model->rope_cos, model->rope_sin, seq_len, num_heads, num_kv_heads, head_dim);

	/* GPU causal attention (f32) */
	if (!iris_metal_causal_attention(model->attn_out, model->q_buf, model->k_buf, model->v_buf, attention_mask, seq_len, num_heads, num_kv_heads, head_dim,
	                                 scale)) {
		qwen3_attention_forward(model, layer, seq_len, attention_mask);
		return;
	}

	/* Output projection on GPU */
	iris_gpu_tensor_t attn = f32_to_bf16_tensor(model->attn_out, seq_len * q_dim);
	if (!attn) {
		qwen3_linear(model->hidden_state, model->attn_out, layer->attn.o_proj_weight, seq_len, q_dim, hidden);
		return;
	}

	iris_gpu_tensor_t out = iris_gpu_linear_bf16_native(attn, layer->attn.o_proj_weight_bf16, seq_len, q_dim, hidden);
	iris_gpu_tensor_free(attn);

	if (!out) {
		qwen3_linear(model->hidden_state, model->attn_out, layer->attn.o_proj_weight, seq_len, q_dim, hidden);
		return;
	}

	bf16_tensor_to_f32(out, model->hidden_state);
	iris_gpu_tensor_free(out);
}

/* GPU-accelerated layer forward */
static void qwen3_layer_forward_bf16(qwen3_model_t *model, qwen3_layer_t *layer, int seq_len, const int *attention_mask) {
	int hidden = model->hidden_size;

	/* Check if we have bf16 weights */
	if (!layer->attn.q_proj_weight_bf16 || !layer->mlp.gate_proj_weight_bf16) {
		qwen3_layer_forward(model, layer, seq_len, attention_mask);
		return;
	}

	/* Save residual */
	memcpy(model->residual, model->hidden_state, seq_len * hidden * sizeof(float));

	/* Pre-attention LayerNorm - CPU (GPU conversion overhead not worth it) */
	qwen3_rms_norm(model->norm_buf, model->hidden_state, layer->input_layernorm_weight, seq_len, hidden, QWEN3_RMS_NORM_EPS);

	/* Self-attention with GPU projections */
	qwen3_attention_forward_bf16(model, layer, seq_len, attention_mask);

	/* Residual connection */
	for (int i = 0; i < seq_len * hidden; i++) {
		model->hidden_state[i] += model->residual[i];
	}

	/* Save residual */
	memcpy(model->residual, model->hidden_state, seq_len * hidden * sizeof(float));

	/* Pre-MLP LayerNorm - CPU (GPU conversion overhead not worth it) */
	qwen3_rms_norm(model->norm_buf, model->hidden_state, layer->post_attention_layernorm_weight, seq_len, hidden, QWEN3_RMS_NORM_EPS);

	/* MLP with GPU */
	qwen3_mlp_forward_bf16(model, layer, seq_len);

	/* Residual connection */
	for (int i = 0; i < seq_len * hidden; i++) {
		model->hidden_state[i] = model->residual[i] + model->mlp_out[i];
	}
}

/* ========================================================================
 * Fully GPU-Resident Forward Pass
 *
 * Keeps hidden state on GPU (bf16) across all layers, eliminating the
 * 72 CPU-GPU syncs (2 per layer × 36 layers) of the mixed path above.
 * Only one GPU sync at the end for all 27 needed layers.
 * ======================================================================== */

/* GPU-only attention: bf16 tensor in, bf16 tensor out.
 * Returns O projection output or NULL on failure. */
static iris_gpu_tensor_t qwen3_attention_gpu(qwen3_model_t *model, qwen3_layer_t *layer, iris_gpu_tensor_t norm_out, int seq_len, const int *attention_mask) {
	int   num_heads    = model->num_heads;
	int   num_kv_heads = model->num_kv_heads;
	int   head_dim     = model->head_dim;
	int   hidden       = model->hidden_size;
	int   q_dim        = num_heads * head_dim;
	int   kv_dim       = num_kv_heads * head_dim;
	float scale        = 1.0f / sqrtf((float)head_dim);

	/* Q, K, V projections (bf16 → bf16) */
	iris_gpu_tensor_t q = iris_gpu_linear_bf16_native(norm_out, layer->attn.q_proj_weight_bf16, seq_len, hidden, q_dim);
	iris_gpu_tensor_t k = iris_gpu_linear_bf16_native(norm_out, layer->attn.k_proj_weight_bf16, seq_len, hidden, kv_dim);
	iris_gpu_tensor_t v = iris_gpu_linear_bf16_native(norm_out, layer->attn.v_proj_weight_bf16, seq_len, hidden, kv_dim);
	if (!q || !k || !v)
		goto fail_qkv;

	/* Q/K RMS normalization on GPU */
	if (layer->attn.q_norm_weight_bf16 && layer->attn.k_norm_weight_bf16) {
		iris_gpu_tensor_t q_norm_w = bf16_ptr_to_bf16_tensor(layer->attn.q_norm_weight_bf16, head_dim);
		iris_gpu_tensor_t k_norm_w = bf16_ptr_to_bf16_tensor(layer->attn.k_norm_weight_bf16, head_dim);
		if (!q_norm_w || !k_norm_w) {
			if (q_norm_w)
				iris_gpu_tensor_free(q_norm_w);
			if (k_norm_w)
				iris_gpu_tensor_free(k_norm_w);
			goto fail_qkv;
		}
		iris_gpu_head_rms_norm_bf16(q, q_norm_w, seq_len, num_heads, head_dim, QWEN3_RMS_NORM_EPS);
		iris_gpu_head_rms_norm_bf16(k, k_norm_w, seq_len, num_kv_heads, head_dim, QWEN3_RMS_NORM_EPS);
		iris_gpu_tensor_free(q_norm_w);
		iris_gpu_tensor_free(k_norm_w);
	}

	/* RoPE */
	iris_gpu_rope_text_bf16(q, k, model->rope_cos, model->rope_sin, seq_len, num_heads, num_kv_heads, head_dim);

	/* Causal attention */
	iris_gpu_tensor_t attn_out = iris_gpu_tensor_alloc_f16(seq_len * q_dim);
	if (!attn_out)
		goto fail_qkv;
	if (!iris_gpu_causal_attention_bf16(attn_out, q, k, v, attention_mask, seq_len, num_heads, num_kv_heads, head_dim, scale)) {
		iris_gpu_tensor_free(attn_out);
		goto fail_qkv;
	}
	iris_gpu_tensor_free(q);
	iris_gpu_tensor_free(k);
	iris_gpu_tensor_free(v);

	/* O projection */
	iris_gpu_tensor_t out = iris_gpu_linear_bf16_native(attn_out, layer->attn.o_proj_weight_bf16, seq_len, q_dim, hidden);
	iris_gpu_tensor_free(attn_out);
	return out;

fail_qkv:
	if (q)
		iris_gpu_tensor_free(q);
	if (k)
		iris_gpu_tensor_free(k);
	if (v)
		iris_gpu_tensor_free(v);
	return NULL;
}

/* GPU-only MLP: bf16 tensor in, bf16 tensor out.
 * Returns down projection output or NULL on failure. */
static iris_gpu_tensor_t qwen3_mlp_gpu(qwen3_model_t *model, qwen3_layer_t *layer, iris_gpu_tensor_t norm_out, int seq_len) {
	int hidden       = model->hidden_size;
	int intermediate = model->intermediate_size;

	iris_gpu_tensor_t gate = iris_gpu_linear_bf16_native(norm_out, layer->mlp.gate_proj_weight_bf16, seq_len, hidden, intermediate);
	iris_gpu_tensor_t up   = iris_gpu_linear_bf16_native(norm_out, layer->mlp.up_proj_weight_bf16, seq_len, hidden, intermediate);
	if (!gate || !up) {
		if (gate)
			iris_gpu_tensor_free(gate);
		if (up)
			iris_gpu_tensor_free(up);
		return NULL;
	}

	iris_gpu_silu_mul_bf16(gate, up, seq_len * intermediate);
	iris_gpu_tensor_free(up);

	iris_gpu_tensor_t out = iris_gpu_linear_bf16_native(gate, layer->mlp.down_proj_weight_bf16, seq_len, intermediate, hidden);
	iris_gpu_tensor_free(gate);
	return out;
}

/* Fully GPU-resident forward pass.
 * hidden_state must already be set (from embedding lookup).
 * Fills model->layer_outputs[0..2] with layers 8, 17, 26 outputs.
 * Processes layers 0..26 and saves outputs at layers 8, 17, 26.
 * Returns 1 on success, 0 on failure. */
/* Fully GPU-resident Qwen3 forward pass. Key optimization: keeps all hidden
 * states on GPU (bf16) across all layers, only reading back the extracted
 * layer outputs at the end. This reduces ~72 CPU-GPU round-trips (2 per
 * layer x 36 layers) down to a single GPU sync. Layer weights are still
 * loaded from mmap on demand and freed after each layer. */
static int qwen3_forward_gpu(qwen3_model_t *model, int seq_len, const int *attention_mask) {
	int hidden     = model->hidden_size;
	int last_layer = QWEN3_OUTPUT_LAYER_3;
	int num_saved  = 3;

	/* Upload hidden state to GPU as bf16 */
	iris_gpu_batch_begin();
	iris_gpu_tensor_t hidden_gpu = f32_to_bf16_tensor(model->hidden_state, seq_len * hidden);
	if (!hidden_gpu) {
		iris_gpu_batch_end();
		return 0;
	}

	/* Allocate tensors for saved layer outputs */
	iris_gpu_tensor_t saved[3] = {NULL, NULL, NULL};
	for (int i = 0; i < num_saved; i++) {
		saved[i] = iris_gpu_tensor_alloc_f16(seq_len * hidden);
		if (!saved[i]) {
			for (int j = 0; j <= i; j++)
				if (saved[j])
					iris_gpu_tensor_free(saved[j]);
			iris_gpu_tensor_free(hidden_gpu);
			iris_gpu_batch_end();
			return 0;
		}
	}

	int ok = 1;

	for (int layer_idx = 0; layer_idx <= last_layer; layer_idx++) {
		qwen3_layer_t *layer = &model->layers[layer_idx];

		/* Load weights on demand (mmap mode) */
		if (model->use_mmap) {
			if (load_layer_weights_small_f32(layer, model->sf_files, model->num_sf_files, layer_idx) != 0) {
				ok = 0;
				break;
			}
			load_layer_weights_bf16(layer, model->sf_files, model->num_sf_files, layer_idx);
		}

		if (!layer->attn.q_proj_weight_bf16 || !layer->mlp.gate_proj_weight_bf16) {
			ok = 0;
			break;
		}

		/* Input RMS norm */
		iris_gpu_tensor_t norm_w   = f32_to_bf16_tensor(layer->input_layernorm_weight, hidden);
		iris_gpu_tensor_t norm_out = iris_gpu_tensor_alloc_f16(seq_len * hidden);
		if (!norm_w || !norm_out) {
			if (norm_w)
				iris_gpu_tensor_free(norm_w);
			if (norm_out)
				iris_gpu_tensor_free(norm_out);
			ok = 0;
			break;
		}
		iris_gpu_rms_norm_bf16(norm_out, hidden_gpu, norm_w, seq_len, hidden, QWEN3_RMS_NORM_EPS);
		iris_gpu_tensor_free(norm_w);

		/* Attention */
		iris_gpu_tensor_t attn_out = qwen3_attention_gpu(model, layer, norm_out, seq_len, attention_mask);
		iris_gpu_tensor_free(norm_out);
		if (!attn_out) {
			ok = 0;
			break;
		}

		/* Residual: hidden += attn_out */
		iris_gpu_add_bf16(hidden_gpu, hidden_gpu, attn_out, seq_len * hidden);
		iris_gpu_tensor_free(attn_out);

		/* Post-attention RMS norm */
		iris_gpu_tensor_t post_norm_w = f32_to_bf16_tensor(layer->post_attention_layernorm_weight, hidden);
		norm_out                      = iris_gpu_tensor_alloc_f16(seq_len * hidden);
		if (!post_norm_w || !norm_out) {
			if (post_norm_w)
				iris_gpu_tensor_free(post_norm_w);
			if (norm_out)
				iris_gpu_tensor_free(norm_out);
			ok = 0;
			break;
		}
		iris_gpu_rms_norm_bf16(norm_out, hidden_gpu, post_norm_w, seq_len, hidden, QWEN3_RMS_NORM_EPS);
		iris_gpu_tensor_free(post_norm_w);

		/* MLP */
		iris_gpu_tensor_t mlp_out = qwen3_mlp_gpu(model, layer, norm_out, seq_len);
		iris_gpu_tensor_free(norm_out);
		if (!mlp_out) {
			ok = 0;
			break;
		}

		/* Residual: hidden += mlp_out */
		iris_gpu_add_bf16(hidden_gpu, hidden_gpu, mlp_out, seq_len * hidden);
		iris_gpu_tensor_free(mlp_out);

		/* Free layer weights (mmap mode) */
		if (model->use_mmap) {
			free_layer_weights(layer);
		}

		/* Save output at extraction layers (GPU blit copy) */
		if (layer_idx == QWEN3_OUTPUT_LAYER_1)
			iris_gpu_copy_bf16(saved[0], hidden_gpu, seq_len * hidden);
		else if (layer_idx == QWEN3_OUTPUT_LAYER_2)
			iris_gpu_copy_bf16(saved[1], hidden_gpu, seq_len * hidden);
		else if (layer_idx == QWEN3_OUTPUT_LAYER_3)
			iris_gpu_copy_bf16(saved[2], hidden_gpu, seq_len * hidden);

		if (iris_text_progress_callback)
			iris_text_progress_callback(layer_idx, model->num_layers);
	}

	if (!ok) {
		iris_gpu_batch_end();
		for (int i = 0; i < num_saved; i++)
			if (saved[i])
				iris_gpu_tensor_free(saved[i]);
		iris_gpu_tensor_free(hidden_gpu);
		return 0;
	}

	/* Convert saved bf16 → f32 on GPU (still within batch) */
	iris_gpu_tensor_t saved_f32[3] = {NULL, NULL, NULL};
	for (int i = 0; i < num_saved; i++)
		saved_f32[i] = iris_gpu_tensor_bf16_to_f32(saved[i]);

	/* Execute everything in one GPU sync */
	iris_gpu_batch_end();

	/* Signal full completion for progress display */
	if (iris_text_progress_callback)
		iris_text_progress_callback(model->num_layers - 1, model->num_layers);

	/* Read f32 results to CPU */
	int read_ok = 1;
	for (int i = 0; i < num_saved; i++) {
		if (saved_f32[i]) {
			iris_gpu_tensor_read(saved_f32[i], model->layer_outputs[i]);
			iris_gpu_tensor_free(saved_f32[i]);
		}
		else {
			read_ok = 0;
		}
		iris_gpu_tensor_free(saved[i]);
	}
	iris_gpu_tensor_free(hidden_gpu);

	return read_ok;
}

#endif /* USE_METAL */

#ifdef USE_VULKAN
/* Dispatch a resident linear, picking the Q8_0-dequantizing GEMM shader for
 * Q8 checkpoints and the bf16 GEMM otherwise. The weight pointer is opaque in
 * both cases (raw mmap bytes uploaded once and cached by pointer). */
static void qwen3_vk_linear(qwen3_model_t *model, iris_gpu_tensor_t out, iris_gpu_tensor_t x, const uint16_t *weight, int seq, int in_dim, int out_dim) {
	if (model->weights_q8)
		iris_vk_qwen_linear_q8(out, x, (const void *)weight, seq, in_dim, out_dim);
	else
		iris_vk_qwen_linear(out, x, weight, seq, in_dim, out_dim);
}

/* ========================================================================
 * Fully GPU-Resident Forward Pass (Vulkan)
 *
 * Mirrors the Metal resident path: activations stay in f32 device tensors
 * across all layers, weights are bf16 applied in-shader, and the whole layer
 * stack records into a single batch that is submitted once. This eliminates
 * the per-op CPU<->GPU round trips of the GEMM-offload path (the cause of the
 * slow "Encoding text" stage). A fixed set of scratch tensors is reused across
 * layers to avoid per-op allocation churn.
 * ======================================================================== */
static int qwen3_forward_vulkan(qwen3_model_t *model, int seq_len, const int *attention_mask) {
	int   hidden       = model->hidden_size;
	int   num_heads    = model->num_heads;
	int   num_kv_heads = model->num_kv_heads;
	int   head_dim     = model->head_dim;
	int   intermediate = model->intermediate_size;
	int   q_dim        = num_heads * head_dim;
	int   kv_dim       = num_kv_heads * head_dim;
	float scale        = 1.0f / sqrtf((float)head_dim);
	int   last_layer   = QWEN3_OUTPUT_LAYER_3;
	int   num_saved    = 3;

	/* The attention shader bounds the key sequence by 512 (shared memory). */
	if (seq_len > QWEN3_MAX_SEQ_LEN)
		return 0;

	/* f32 attention mask: 1.0 for valid tokens, 0.0 for padding. */
	float *mask_f32 = malloc((size_t)seq_len * sizeof(float));
	if (!mask_f32)
		return 0;
	for (int i = 0; i < seq_len; i++)
		mask_f32[i] = (attention_mask && attention_mask[i] == 0) ? 0.0f : 1.0f;

	iris_gpu_batch_begin();

	iris_gpu_tensor_t hidden_t = iris_gpu_tensor_create(model->hidden_state, (size_t)seq_len * hidden);
	iris_gpu_tensor_t mask_t   = iris_gpu_tensor_create(mask_f32, (size_t)seq_len);
	free(mask_f32);

	/* Reused scratch (allocated once, written each layer). */
	iris_gpu_tensor_t norm_t   = iris_gpu_tensor_alloc((size_t)seq_len * hidden);
	iris_gpu_tensor_t q_t      = iris_gpu_tensor_alloc((size_t)seq_len * q_dim);
	iris_gpu_tensor_t k_t      = iris_gpu_tensor_alloc((size_t)seq_len * kv_dim);
	iris_gpu_tensor_t v_t      = iris_gpu_tensor_alloc((size_t)seq_len * kv_dim);
	iris_gpu_tensor_t attn_t   = iris_gpu_tensor_alloc((size_t)seq_len * q_dim);
	iris_gpu_tensor_t gate_t   = iris_gpu_tensor_alloc((size_t)seq_len * intermediate);
	iris_gpu_tensor_t up_t     = iris_gpu_tensor_alloc((size_t)seq_len * intermediate);
	iris_gpu_tensor_t proj_t   = iris_gpu_tensor_alloc((size_t)seq_len * hidden); /* o / down output */
	iris_gpu_tensor_t saved[3] = {NULL, NULL, NULL};
	for (int i = 0; i < num_saved; i++)
		saved[i] = iris_gpu_tensor_alloc((size_t)seq_len * hidden);

	int ok = hidden_t && mask_t && norm_t && q_t && k_t && v_t && attn_t && gate_t && up_t && proj_t;
	for (int i = 0; i < num_saved; i++)
		if (!saved[i])
			ok = 0;

	for (int layer_idx = 0; ok && layer_idx <= last_layer; layer_idx++) {
		qwen3_layer_t *layer = &model->layers[layer_idx];
		if (model->use_mmap)
			load_layer_weights_bf16(layer, model->sf_files, model->num_sf_files, layer_idx);

		qwen3_attention_t *a   = &layer->attn;
		qwen3_mlp_t       *mlp = &layer->mlp;
		if (!a->q_proj_weight_bf16 || !a->k_proj_weight_bf16 || !a->v_proj_weight_bf16 || !a->o_proj_weight_bf16 || !a->q_norm_weight_bf16 ||
		    !a->k_norm_weight_bf16 || !mlp->gate_proj_weight_bf16 || !mlp->up_proj_weight_bf16 || !mlp->down_proj_weight_bf16 ||
		    !layer->input_layernorm_weight_bf16 || !layer->post_attention_layernorm_weight_bf16) {
			ok = 0;
			if (model->use_mmap)
				free_layer_weights(layer);
			break;
		}

		/* Self-attention */
		iris_vk_qwen_rms_norm(norm_t, hidden_t, layer->input_layernorm_weight_bf16, seq_len, hidden, QWEN3_RMS_NORM_EPS);
		qwen3_vk_linear(model, q_t, norm_t, a->q_proj_weight_bf16, seq_len, hidden, q_dim);
		qwen3_vk_linear(model, k_t, norm_t, a->k_proj_weight_bf16, seq_len, hidden, kv_dim);
		qwen3_vk_linear(model, v_t, norm_t, a->v_proj_weight_bf16, seq_len, hidden, kv_dim);
		iris_vk_qwen_head_rms_norm(q_t, a->q_norm_weight_bf16, seq_len, num_heads, head_dim, QWEN3_RMS_NORM_EPS);
		iris_vk_qwen_head_rms_norm(k_t, a->k_norm_weight_bf16, seq_len, num_kv_heads, head_dim, QWEN3_RMS_NORM_EPS);
		iris_vk_qwen_rope(q_t, k_t, model->rope_cos, model->rope_sin, seq_len, num_heads, num_kv_heads, head_dim);
		iris_vk_qwen_attention(attn_t, q_t, k_t, v_t, mask_t, seq_len, num_heads, num_kv_heads, head_dim, scale);
		qwen3_vk_linear(model, proj_t, attn_t, a->o_proj_weight_bf16, seq_len, q_dim, hidden);
		iris_gpu_add_f32(hidden_t, hidden_t, proj_t, seq_len * hidden);

		/* MLP (SwiGLU) */
		iris_vk_qwen_rms_norm(norm_t, hidden_t, layer->post_attention_layernorm_weight_bf16, seq_len, hidden, QWEN3_RMS_NORM_EPS);
		qwen3_vk_linear(model, gate_t, norm_t, mlp->gate_proj_weight_bf16, seq_len, hidden, intermediate);
		qwen3_vk_linear(model, up_t, norm_t, mlp->up_proj_weight_bf16, seq_len, hidden, intermediate);
		iris_vk_qwen_silu_mul(gate_t, up_t, seq_len * intermediate);
		qwen3_vk_linear(model, proj_t, gate_t, mlp->down_proj_weight_bf16, seq_len, intermediate, hidden);
		iris_gpu_add_f32(hidden_t, hidden_t, proj_t, seq_len * hidden);

		if (model->use_mmap)
			free_layer_weights(layer);

		/* Save output at extraction layers (GPU-side copy). */
		if (layer_idx == QWEN3_OUTPUT_LAYER_1)
			iris_gpu_copy_f32(saved[0], hidden_t, (size_t)seq_len * hidden);
		else if (layer_idx == QWEN3_OUTPUT_LAYER_2)
			iris_gpu_copy_f32(saved[1], hidden_t, (size_t)seq_len * hidden);
		else if (layer_idx == QWEN3_OUTPUT_LAYER_3)
			iris_gpu_copy_f32(saved[2], hidden_t, (size_t)seq_len * hidden);

		if (iris_text_progress_callback)
			iris_text_progress_callback(layer_idx, model->num_layers);
	}

	/* Execute the whole stack in one submit, then read extracted outputs. */
	iris_gpu_batch_end();

	if (ok) {
		if (iris_text_progress_callback)
			iris_text_progress_callback(model->num_layers - 1, model->num_layers);
		for (int i = 0; i < num_saved; i++)
			iris_gpu_tensor_read(saved[i], model->layer_outputs[i]);
	}

	iris_gpu_tensor_free(hidden_t);
	iris_gpu_tensor_free(mask_t);
	iris_gpu_tensor_free(norm_t);
	iris_gpu_tensor_free(q_t);
	iris_gpu_tensor_free(k_t);
	iris_gpu_tensor_free(v_t);
	iris_gpu_tensor_free(attn_t);
	iris_gpu_tensor_free(gate_t);
	iris_gpu_tensor_free(up_t);
	iris_gpu_tensor_free(proj_t);
	for (int i = 0; i < num_saved; i++)
		iris_gpu_tensor_free(saved[i]);

	return ok;
}
#endif /* USE_VULKAN */

/* ========================================================================
 * Forward Pass
 * ======================================================================== */

/* Main Qwen3 forward pass. Runs embedding lookup then processes through
 * transformer layers, saving hidden states at layers 8, 17, 26 and
 * concatenating them -> [seq, 3*hidden]. Stops early at the last needed
 * extraction layer to skip ~9 unnecessary layers of compute. Tries fully
 * GPU-resident path first, falls back to mixed CPU/GPU or pure CPU. */
float *qwen3_forward(qwen3_model_t *model, const int *input_ids, const int *attention_mask, int seq_len) {
	int    hidden     = model->hidden_size;
	int    last_layer = QWEN3_OUTPUT_LAYER_3;
	float *output;

	/* Embedding lookup */
	for (int s = 0; s < seq_len; s++) {
		int token_id = input_ids[s];
		if (token_id >= 0 && token_id < model->vocab_size) {
			memcpy(model->hidden_state + s * hidden, model->embed_tokens + token_id * hidden, hidden * sizeof(float));
		}
		else {
			/* Unknown token - use zeros */
			memset(model->hidden_state + s * hidden, 0, hidden * sizeof(float));
		}
	}

	/* Run through transformer layers */
#ifdef USE_VULKAN
	/* Try the fully GPU-resident path: one submit for the whole layer stack
	 * instead of a CPU<->GPU round trip per linear op. */
	if (model->use_bf16 && iris_vk_qwen_available() && seq_len <= QWEN3_MAX_SEQ_LEN) {
		if (qwen3_forward_vulkan(model, seq_len, attention_mask))
			goto concatenate;
		/* Resident path failed; hidden_state is untouched (GPU worked on a
		 * copy). Fall through to the GEMM-offload CPU/GPU path below. */
	}
#endif
#ifdef USE_METAL
	/* Try fully GPU-resident path: 1 sync instead of 72, skips unneeded layers */
	if (model->use_bf16 && iris_metal_available() && seq_len <= 512) {
		if (qwen3_forward_gpu(model, seq_len, attention_mask))
			goto concatenate;
		/* GPU path failed, fall through to mixed CPU/GPU path.
		 * hidden_state is unmodified (GPU worked on a copy). */
	}

	/* Start batch mode to reduce GPU sync overhead between layers */
	int batch_mode = model->use_bf16 && iris_metal_available();
	if (batch_mode) {
		iris_gpu_batch_begin();
	}
#endif

	for (int layer_idx = 0; layer_idx <= last_layer; layer_idx++) {
		/* In mmap mode, load layer weights on-demand */
		if (model->use_mmap) {
#ifdef USE_METAL
			if (model->use_bf16) {
				/* Load only small f32 weights (layer norms) + bf16 projection weights */
				if (load_layer_weights_small_f32(&model->layers[layer_idx], model->sf_files, model->num_sf_files, layer_idx) != 0) {
					fprintf(stderr, "Failed to load layer %d small weights\n", layer_idx);
#ifdef USE_METAL
					if (batch_mode)
						iris_gpu_batch_end();
#endif
					return NULL;
				}
				load_layer_weights_bf16(&model->layers[layer_idx], model->sf_files, model->num_sf_files, layer_idx);
			}
			else
#endif
			{
				if (load_layer_weights(&model->layers[layer_idx], model->sf_files, model->num_sf_files, layer_idx) != 0) {
					fprintf(stderr, "Failed to load layer %d weights\n", layer_idx);
					return NULL;
				}
			}
		}

#ifdef USE_METAL
		if (model->use_bf16 && iris_metal_available()) {
			qwen3_layer_forward_bf16(model, &model->layers[layer_idx], seq_len, attention_mask);
		}
		else
#endif
		{
			qwen3_layer_forward(model, &model->layers[layer_idx], seq_len, attention_mask);
		}

		/* In mmap mode, free layer weights after use */
		if (model->use_mmap) {
			free_layer_weights(&model->layers[layer_idx]);
		}

		/* Save output at extraction layers */
		if (layer_idx == QWEN3_OUTPUT_LAYER_1)
			memcpy(model->layer_outputs[0], model->hidden_state, seq_len * hidden * sizeof(float));
		else if (layer_idx == QWEN3_OUTPUT_LAYER_2)
			memcpy(model->layer_outputs[1], model->hidden_state, seq_len * hidden * sizeof(float));
		else if (layer_idx == QWEN3_OUTPUT_LAYER_3)
			memcpy(model->layer_outputs[2], model->hidden_state, seq_len * hidden * sizeof(float));

		/* Progress callback */
		if (iris_text_progress_callback)
			iris_text_progress_callback(layer_idx, model->num_layers);
	}

#ifdef USE_METAL
	/* End batch mode */
	if (batch_mode) {
		iris_gpu_batch_end();
	}
#endif

	/* Build output embeddings */
#if defined(USE_METAL) || defined(USE_VULKAN)
concatenate:
	(void)0; /* label needs a statement; can't precede a declaration in C */
#endif
	int text_dim = model->text_dim;
	output       = malloc(seq_len * text_dim * sizeof(float));
	if (!output)
		return NULL;

	/* Flux: concatenate layers 8, 17, 26 -> [seq_len, 3*hidden] */
	for (int s = 0; s < seq_len; s++) {
		memcpy(output + s * text_dim, model->layer_outputs[0] + s * hidden, hidden * sizeof(float));
		memcpy(output + s * text_dim + hidden, model->layer_outputs[1] + s * hidden, hidden * sizeof(float));
		memcpy(output + s * text_dim + 2 * hidden, model->layer_outputs[2] + s * hidden, hidden * sizeof(float));
	}

	return output;
}

/* ========================================================================
 * Model Loading
 * ======================================================================== */

/* Translate an HF/safetensors Qwen3 tensor name to its llama.cpp GGUF
 * equivalent (e.g. "model.layers.5.self_attn.q_proj.weight" -> "blk.5.attn_q.weight").
 * BF16 GGUF payloads are byte-identical to safetensors and Qwen3 uses NEOX-style
 * RoPE, so no weight permutation is needed - only the names differ. Returns out
 * on success, or NULL if the name has no mapping (caller falls back to the
 * original). */
static const char *gguf_qwen3_name(const char *hf, char *out, size_t outsz) {
	if (strcmp(hf, "model.embed_tokens.weight") == 0) {
		snprintf(out, outsz, "token_embd.weight");
		return out;
	}
	if (strcmp(hf, "model.norm.weight") == 0) {
		snprintf(out, outsz, "output_norm.weight");
		return out;
	}

	int  layer = -1;
	char suffix[128];
	if (sscanf(hf, "model.layers.%d.%127s", &layer, suffix) == 2) {
		const char *g = NULL;
		if (strcmp(suffix, "input_layernorm.weight") == 0)
			g = "attn_norm";
		else if (strcmp(suffix, "post_attention_layernorm.weight") == 0)
			g = "ffn_norm";
		else if (strcmp(suffix, "self_attn.q_proj.weight") == 0)
			g = "attn_q";
		else if (strcmp(suffix, "self_attn.k_proj.weight") == 0)
			g = "attn_k";
		else if (strcmp(suffix, "self_attn.v_proj.weight") == 0)
			g = "attn_v";
		else if (strcmp(suffix, "self_attn.o_proj.weight") == 0)
			g = "attn_output";
		else if (strcmp(suffix, "self_attn.q_norm.weight") == 0)
			g = "attn_q_norm";
		else if (strcmp(suffix, "self_attn.k_norm.weight") == 0)
			g = "attn_k_norm";
		else if (strcmp(suffix, "mlp.gate_proj.weight") == 0)
			g = "ffn_gate";
		else if (strcmp(suffix, "mlp.up_proj.weight") == 0)
			g = "ffn_up";
		else if (strcmp(suffix, "mlp.down_proj.weight") == 0)
			g = "ffn_down";
		if (g) {
			snprintf(out, outsz, "blk.%d.%s.weight", layer, g);
			return out;
		}
	}
	return NULL;
}

/* Resolve a lookup name for a given file: GGUF files (header_json == NULL) use
 * llama.cpp tensor names, so translate the HF name on the fly. */
static const char *qwen3_lookup_name(const safetensors_file_t *sf, const char *name, char *buf, size_t bufsz) {
	if (sf->header_json == NULL) {
		const char *g = gguf_qwen3_name(name, buf, bufsz);
		if (g)
			return g;
	}
	return name;
}

/* Helper to load a tensor from safetensors/GGUF files */
static float *load_tensor(safetensors_file_t **files, int num_files, const char *name) {
	char buf[256];
	for (int f = 0; f < num_files; f++) {
		const char         *lookup = qwen3_lookup_name(files[f], name, buf, sizeof(buf));
		const safetensor_t *t      = safetensors_find(files[f], lookup);
		if (t) {
			return safetensors_get_f32(files[f], t);
		}
	}
	fprintf(stderr, "Error: required tensor not found: %s\n", name);
	return NULL;
}

#if defined(USE_METAL) || defined(USE_VULKAN)
/* Load a large GPU weight directly (zero-copy from the mmap region).
 *
 * Returns a raw pointer into the mmap'd file for either a bf16 weight or a
 * GGML Q8_0 block-quantized weight - in both cases the bytes are uploaded to
 * VRAM verbatim and dequantized in-shader, so the pointer type is opaque (the
 * model's weights_q8 flag tells the resident path which GEMM shader to use).
 * Returns NULL if the tensor is absent or in an unsupported format. */
static uint16_t *load_tensor_bf16(safetensors_file_t **files, int num_files, const char *name) {
	char buf[256];
	for (int f = 0; f < num_files; f++) {
		const char         *lookup = qwen3_lookup_name(files[f], name, buf, sizeof(buf));
		const safetensor_t *t      = safetensors_find(files[f], lookup);
		if (t && safetensor_is_bf16(t)) {
			return safetensors_get_bf16_direct(files[f], t);
		}
		if (t && safetensor_is_q8_0(t)) {
			uint16_t *p = (uint16_t *)safetensors_data(files[f], t);
#ifdef USE_METAL
			/* Metal consumes Q8_0 through the bf16 weight APIs and dequantizes
			 * at upload; flag the pointer so the backend knows to do so.
			 * (Vulkan dequantizes Q8_0 in-shader and needs no registration.) */
			iris_metal_register_q8_weight(p);
#endif
			return p;
		}
	}
	return NULL; /* Not found or unsupported format - fall back to f32 */
}

/* Load a tensor as an owned bf16 buffer (caller must free).
 *
 * The GPU-resident text-encoder path needs every weight - including the RMS
 * norm weights - as bf16. HF safetensors stores the norms as bf16, but a
 * llama.cpp GGUF stores all *.norm weights as F32, so the zero-copy
 * load_tensor_bf16() returns NULL for them and the resident path bails out to a
 * ~15x slower CPU fallback. Here we copy bf16 sources directly and truncate F32
 * sources to bf16 (the GGUF F32 norm was upcast from bf16, so truncation is
 * lossless), giving the resident path the bf16 norms it expects regardless of
 * container. */
static uint16_t *load_tensor_bf16_owned(safetensors_file_t **files, int num_files, const char *name) {
	char buf[256];
	for (int f = 0; f < num_files; f++) {
		const char         *lookup = qwen3_lookup_name(files[f], name, buf, sizeof(buf));
		const safetensor_t *t      = safetensors_find(files[f], lookup);
		if (!t)
			continue;
		if (safetensor_is_bf16(t)) {
			return safetensors_get_bf16(files[f], t); /* allocates + copies */
		}
		/* F32 (or F16-as-f32) source: dequantize then truncate to bf16. */
		float *f32 = safetensors_get_f32(files[f], t);
		if (!f32)
			return NULL;
		int64_t   n   = safetensor_numel(t);
		uint16_t *out = malloc((size_t)n * sizeof(uint16_t));
		if (out) {
			for (int64_t i = 0; i < n; i++) {
				uint32_t bits;
				memcpy(&bits, &f32[i], sizeof(bits));
				out[i] = (uint16_t)(bits >> 16);
			}
		}
		free(f32);
		return out;
	}
	return NULL;
}
#endif

#ifdef USE_METAL
/* Load only small f32 weights for bf16 path (layer norms and q/k norms) */
static int load_layer_weights_small_f32(qwen3_layer_t *layer, safetensors_file_t **files, int num_files, int layer_idx) {
	char name[256];

	/* Read in GGUF on-disk order: attn_norm, ffn_norm, attn_k_norm, attn_q_norm. */

	/* Input layernorm */
	snprintf(name, sizeof(name), "model.layers.%d.input_layernorm.weight", layer_idx);
	layer->input_layernorm_weight = load_tensor(files, num_files, name);

	/* Post-attention layernorm */
	snprintf(name, sizeof(name), "model.layers.%d.post_attention_layernorm.weight", layer_idx);
	layer->post_attention_layernorm_weight = load_tensor(files, num_files, name);

	/* K/Q norm */
	snprintf(name, sizeof(name), "model.layers.%d.self_attn.k_norm.weight", layer_idx);
	layer->attn.k_norm_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.q_norm.weight", layer_idx);
	layer->attn.q_norm_weight = load_tensor(files, num_files, name);

	return (layer->input_layernorm_weight && layer->post_attention_layernorm_weight && layer->attn.q_norm_weight && layer->attn.k_norm_weight) ? 0 : -1;
}
#endif

static int load_layer_weights(qwen3_layer_t *layer, safetensors_file_t **files, int num_files, int layer_idx) {
	char name[256];

	/* Read tensors in the order they are laid out in the GGUF file
	 * (attn_norm, ffn_down/gate/up, ffn_norm, attn_k_norm/k, attn_output,
	 * attn_q_norm/q, attn_v) rather than HF/forward order. In mmap mode each
	 * load_tensor() faults the tensor's pages straight off disk, so matching the
	 * on-disk layout keeps those reads sequential and lets kernel readahead
	 * work, which is what was making "Encoding text..." slow for GGUF. */

	/* Input layernorm */
	snprintf(name, sizeof(name), "model.layers.%d.input_layernorm.weight", layer_idx);
	layer->input_layernorm_weight = load_tensor(files, num_files, name);

	/* MLP weights */
	snprintf(name, sizeof(name), "model.layers.%d.mlp.down_proj.weight", layer_idx);
	layer->mlp.down_proj_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.mlp.gate_proj.weight", layer_idx);
	layer->mlp.gate_proj_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.mlp.up_proj.weight", layer_idx);
	layer->mlp.up_proj_weight = load_tensor(files, num_files, name);

	/* Post-attention layernorm */
	snprintf(name, sizeof(name), "model.layers.%d.post_attention_layernorm.weight", layer_idx);
	layer->post_attention_layernorm_weight = load_tensor(files, num_files, name);

	/* K norm + attention K/output weights */
	snprintf(name, sizeof(name), "model.layers.%d.self_attn.k_norm.weight", layer_idx);
	layer->attn.k_norm_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.k_proj.weight", layer_idx);
	layer->attn.k_proj_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.o_proj.weight", layer_idx);
	layer->attn.o_proj_weight = load_tensor(files, num_files, name);

	/* Q norm + attention Q/V weights */
	snprintf(name, sizeof(name), "model.layers.%d.self_attn.q_norm.weight", layer_idx);
	layer->attn.q_norm_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.q_proj.weight", layer_idx);
	layer->attn.q_proj_weight = load_tensor(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.v_proj.weight", layer_idx);
	layer->attn.v_proj_weight = load_tensor(files, num_files, name);

	/* Check that all required tensors were loaded */
	if (!layer->input_layernorm_weight || !layer->post_attention_layernorm_weight || !layer->attn.q_proj_weight || !layer->attn.k_proj_weight ||
	    !layer->attn.v_proj_weight || !layer->attn.o_proj_weight || !layer->attn.q_norm_weight || !layer->attn.k_norm_weight || !layer->mlp.gate_proj_weight ||
	    !layer->mlp.up_proj_weight || !layer->mlp.down_proj_weight) {
		return -1;
	}

	return 0;
}

#if defined(USE_METAL) || defined(USE_VULKAN)
/* Load bf16 weights for a layer (GPU acceleration path).
 * Returns 1 if all bf16 weights loaded successfully, 0 otherwise.
 * bf16 pointers are direct into mmap region - do NOT free them. */
static int load_layer_weights_bf16(qwen3_layer_t *layer, safetensors_file_t **files, int num_files, int layer_idx) {
	char name[256];

	/* Read tensors in GGUF on-disk order (see load_layer_weights) so the page
	 * faults that back these mmap pointers stay sequential during streaming. */

	/* Norm weights are bf16 in HF safetensors but F32 in a llama.cpp GGUF, so
	 * for GGUF we convert them to bf16. The GPU weight cache keys on the host
	 * pointer and the resident forward is batched, so a converted buffer must
	 * keep a stable, unique address for the whole batch - a per-forward
	 * malloc/free would alias a still-pending dispatch and corrupt the output.
	 * We therefore convert each norm once and keep it on the layer (preserved
	 * across free_layer_weights, released in qwen3_model_free). */
	snprintf(name, sizeof(name), "model.layers.%d.input_layernorm.weight", layer_idx);
	if (!layer->input_layernorm_weight_bf16)
		layer->input_layernorm_weight_bf16 = load_tensor_bf16_owned(files, num_files, name);

	/* MLP weights */
	snprintf(name, sizeof(name), "model.layers.%d.mlp.down_proj.weight", layer_idx);
	layer->mlp.down_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.mlp.gate_proj.weight", layer_idx);
	layer->mlp.gate_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.mlp.up_proj.weight", layer_idx);
	layer->mlp.up_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	/* Post-attention layernorm (converted once; see note above) */
	snprintf(name, sizeof(name), "model.layers.%d.post_attention_layernorm.weight", layer_idx);
	if (!layer->post_attention_layernorm_weight_bf16)
		layer->post_attention_layernorm_weight_bf16 = load_tensor_bf16_owned(files, num_files, name);

	/* K norm (converted once) + attention K/output weights */
	snprintf(name, sizeof(name), "model.layers.%d.self_attn.k_norm.weight", layer_idx);
	if (!layer->attn.k_norm_weight_bf16)
		layer->attn.k_norm_weight_bf16 = load_tensor_bf16_owned(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.k_proj.weight", layer_idx);
	layer->attn.k_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.o_proj.weight", layer_idx);
	layer->attn.o_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	/* Q norm (converted once) + attention Q/V weights */
	snprintf(name, sizeof(name), "model.layers.%d.self_attn.q_norm.weight", layer_idx);
	if (!layer->attn.q_norm_weight_bf16)
		layer->attn.q_norm_weight_bf16 = load_tensor_bf16_owned(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.q_proj.weight", layer_idx);
	layer->attn.q_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	snprintf(name, sizeof(name), "model.layers.%d.self_attn.v_proj.weight", layer_idx);
	layer->attn.v_proj_weight_bf16 = load_tensor_bf16(files, num_files, name);

	/* Check if all large weights loaded as bf16 */
	return (layer->attn.q_proj_weight_bf16 && layer->attn.k_proj_weight_bf16 && layer->attn.v_proj_weight_bf16 && layer->attn.o_proj_weight_bf16 &&
	        layer->mlp.gate_proj_weight_bf16 && layer->mlp.up_proj_weight_bf16 && layer->mlp.down_proj_weight_bf16);
}
#endif

/* Free a single layer's weights (used in mmap streaming mode) */
static void free_layer_weights(qwen3_layer_t *layer) {
	free(layer->input_layernorm_weight);
	free(layer->post_attention_layernorm_weight);
	free(layer->attn.q_proj_weight);
	free(layer->attn.k_proj_weight);
	free(layer->attn.v_proj_weight);
	free(layer->attn.o_proj_weight);
	free(layer->attn.q_norm_weight);
	free(layer->attn.k_norm_weight);
	free(layer->mlp.gate_proj_weight);
	free(layer->mlp.up_proj_weight);
	free(layer->mlp.down_proj_weight);
#if defined(USE_METAL) || defined(USE_VULKAN)
	/* The bf16 norm weights are persistent owned copies that must keep a stable
	 * address across forwards (the GPU weight cache keys on the pointer); carry
	 * them across the reset and release them only in qwen3_model_free. The bf16
	 * projection weights are direct mmap views and need no free. */
	uint16_t *in_norm   = layer->input_layernorm_weight_bf16;
	uint16_t *post_norm = layer->post_attention_layernorm_weight_bf16;
	uint16_t *q_norm    = layer->attn.q_norm_weight_bf16;
	uint16_t *k_norm    = layer->attn.k_norm_weight_bf16;
	memset(layer, 0, sizeof(*layer));
	layer->input_layernorm_weight_bf16          = in_norm;
	layer->post_attention_layernorm_weight_bf16 = post_norm;
	layer->attn.q_norm_weight_bf16              = q_norm;
	layer->attn.k_norm_weight_bf16              = k_norm;
#else
	memset(layer, 0, sizeof(*layer));
#endif
}

/* ========================================================================
 * Qwen3 Config Parsing and Dynamic Shard Loading
 * ======================================================================== */

/* Parse text_encoder/config.json to get architecture dimensions.
 * Sets model fields. Returns 0 on success, -1 on failure. */
static int parse_qwen3_config(const char *model_dir, qwen3_model_t *model) {
	char path[1024];
	snprintf(path, sizeof(path), "%s/config.json", model_dir);

	FILE *f = fopen(path, "r");
	if (!f)
		return -1;

	char   buf[8192];
	size_t n = fread(buf, 1, sizeof(buf) - 1, f);
	buf[n]   = '\0';
	fclose(f);

	char *p;
	int   hidden = 0, intermediate = 0, num_heads = 0, num_kv_heads = 0;
	int   head_dim = 0, vocab_size = 0, num_layers = 0;
	float rope_theta = 0;

	if ((p = strstr(buf, "\"hidden_size\""))) {
		if ((p = strchr(p, ':')))
			hidden = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"intermediate_size\""))) {
		if ((p = strchr(p, ':')))
			intermediate = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"num_attention_heads\""))) {
		if ((p = strchr(p, ':')))
			num_heads = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"num_key_value_heads\""))) {
		if ((p = strchr(p, ':')))
			num_kv_heads = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"head_dim\""))) {
		if ((p = strchr(p, ':')))
			head_dim = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"vocab_size\""))) {
		if ((p = strchr(p, ':')))
			vocab_size = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"num_hidden_layers\""))) {
		if ((p = strchr(p, ':')))
			num_layers = atoi(p + 1);
	}
	if ((p = strstr(buf, "\"rope_theta\""))) {
		if ((p = strchr(p, ':')))
			rope_theta = atof(p + 1);
	}

	if (hidden <= 0 || num_heads <= 0)
		return -1;

	model->hidden_size       = hidden;
	model->intermediate_size = intermediate > 0 ? intermediate : 9728;
	model->num_heads         = num_heads;
	model->num_kv_heads      = num_kv_heads > 0 ? num_kv_heads : 8;
	model->head_dim          = head_dim > 0 ? head_dim : 128;
	model->vocab_size        = vocab_size > 0 ? vocab_size : QWEN3_VOCAB_SIZE;
	model->num_layers        = num_layers > 0 ? num_layers : 36;
	model->rope_theta        = rope_theta > 0 ? rope_theta : QWEN3_ROPE_THETA;
	model->text_dim          = 3 * hidden;

	return 0;
}

/* Set default Qwen3-4B architecture values. */
static void qwen3_set_defaults(qwen3_model_t *model) {
	model->hidden_size       = 2560;
	model->intermediate_size = 9728;
	model->num_heads         = 32;
	model->num_kv_heads      = 8;
	model->head_dim          = 128;
	model->vocab_size        = QWEN3_VOCAB_SIZE;
	model->num_layers        = 36;
	model->rope_theta        = QWEN3_ROPE_THETA;
	model->text_dim          = 3 * 2560;
}

/* Open all safetensors shards dynamically by reading model.safetensors.index.json.
 * Returns number of files opened, 0 on failure. */
static int open_safetensors_shards(const char *model_dir, safetensors_file_t **files, int max_files) {
	char path[1024];

	/* Prefer a GGUF checkpoint if one is present in the model directory.
	 * gguf_open yields a safetensors_file_t with header_json == NULL, which the
	 * tensor loaders detect to translate HF tensor names to llama.cpp GGUF names.
	 * BF16 GGUF payloads are byte-identical to safetensors, so a single file
	 * replaces the safetensors shards entirely. */
	{
		/* Hardcoded text-encoder GGUF filename (lives in the model root). */
		const char *gguf_name = "Qwen3-4B-Q8_0.gguf";
		snprintf(path, sizeof(path), "%s/%s", model_dir, gguf_name);
		FILE *gf = fopen(path, "rb");
		if (gf)
			fclose(gf);
		if (gf && max_files > 0) {
			files[0] = gguf_open(path);
			if (files[0]) {
				if (iris_verbose)
					fprintf(stderr, "qwen3: loading text encoder from GGUF: %s\n", gguf_name);
				return 1;
			}
			fprintf(stderr, "qwen3: failed to open GGUF %s\n", gguf_name);
			return 0;
		}
	}

	/* First try: read the index JSON to discover shard filenames */
	snprintf(path, sizeof(path), "%s/model.safetensors.index.json", model_dir);
	FILE *f = fopen(path, "r");
	if (f) {
		/* Read the whole index file */
		fseek(f, 0, SEEK_END);
		long fsize = ftell(f);
		fseek(f, 0, SEEK_SET);
		char *buf = malloc(fsize + 1);
		if (!buf) {
			fclose(f);
			return 0;
		}
		fread(buf, 1, fsize, f);
		buf[fsize] = '\0';
		fclose(f);

		/* Collect unique shard filenames from weight_map values.
		 * They look like: "model-00001-of-00003.safetensors" */
		char shard_names[QWEN3_MAX_SHARDS][128];
		int  num_shards = 0;

		char *p = buf;
		while ((p = strstr(p, "model-")) != NULL) {
			/* Find the end of the filename */
			char *end = strstr(p, ".safetensors");
			if (!end) {
				p++;
				continue;
			}
			end += strlen(".safetensors");

			int len = (int)(end - p);
			if (len >= 128) {
				p = end;
				continue;
			}

			/* Check if we already have this shard */
			char name[128];
			memcpy(name, p, len);
			name[len] = '\0';

			int found = 0;
			for (int i = 0; i < num_shards; i++) {
				if (strcmp(shard_names[i], name) == 0) {
					found = 1;
					break;
				}
			}
			if (!found && num_shards < max_files && num_shards < QWEN3_MAX_SHARDS) {
				strcpy(shard_names[num_shards], name);
				num_shards++;
			}
			p = end;
		}
		free(buf);

		if (num_shards > 0) {
			for (int i = 0; i < num_shards; i++) {
				snprintf(path, sizeof(path), "%s/%s", model_dir, shard_names[i]);
				files[i] = safetensors_open(path);
				if (!files[i]) {
					fprintf(stderr, "qwen3: failed to open shard %s\n", path);
					for (int j = 0; j < i; j++)
						safetensors_close(files[j]);
					return 0;
				}
			}
			return num_shards;
		}
	}

	/* Fallback: try model-00001-of-00002.safetensors pattern */
	char path1[512], path2[512];
	snprintf(path1, sizeof(path1), "%s/model-00001-of-00002.safetensors", model_dir);
	snprintf(path2, sizeof(path2), "%s/model-00002-of-00002.safetensors", model_dir);

	files[0] = safetensors_open(path1);
	files[1] = safetensors_open(path2);
	if (files[0] && files[1])
		return 2;

	if (files[0])
		safetensors_close(files[0]);
	if (files[1])
		safetensors_close(files[1]);
	return 0;
}

/* Allocate working memory for Qwen3 model based on architecture fields. */
static void qwen3_alloc_work_buffers(qwen3_model_t *model) {
	int seq_len      = QWEN3_MAX_SEQ_LEN;
	int hidden       = model->hidden_size;
	int num_heads    = model->num_heads;
	int num_kv_heads = model->num_kv_heads;
	int head_dim     = model->head_dim;
	int intermediate = model->intermediate_size;

	model->hidden_state = malloc(seq_len * hidden * sizeof(float));
	model->residual     = malloc(seq_len * hidden * sizeof(float));
	model->q_buf        = malloc(seq_len * num_heads * head_dim * sizeof(float));
	model->k_buf        = malloc(seq_len * num_kv_heads * head_dim * sizeof(float));
	model->v_buf        = malloc(seq_len * num_kv_heads * head_dim * sizeof(float));
	model->attn_scores  = malloc(num_heads * seq_len * seq_len * sizeof(float));
	model->attn_out     = malloc(seq_len * num_heads * head_dim * sizeof(float));
	model->mlp_gate     = malloc(seq_len * intermediate * sizeof(float));
	model->mlp_up       = malloc(seq_len * intermediate * sizeof(float));
	model->mlp_out      = malloc(seq_len * hidden * sizeof(float));
	model->norm_buf     = malloc(seq_len * hidden * sizeof(float));

	model->attn_q_head   = malloc(seq_len * head_dim * sizeof(float));
	model->attn_v_head   = malloc(seq_len * head_dim * sizeof(float));
	model->attn_out_head = malloc(seq_len * head_dim * sizeof(float));

	for (int i = 0; i < 3; i++) {
		model->layer_outputs[i] = malloc(seq_len * hidden * sizeof(float));
	}
}

/* Eager-mode model loading: reads all weights into RAM upfront. Parses
 * config.json for architecture parameters (hidden size, head counts, etc.),
 * loads all layer weights from safetensors shards, and precomputes RoPE
 * frequency tables. Uses ~16GB RAM for 4B model. Prefer mmap mode for
 * lower memory usage unless weights need repeated random access. */
qwen3_model_t *qwen3_model_load(const char *model_dir) {
	qwen3_model_t *model = calloc(1, sizeof(qwen3_model_t));
	if (!model)
		return NULL;

	/* Parse config, fall back to defaults */
	if (parse_qwen3_config(model_dir, model) != 0) {
		qwen3_set_defaults(model);
	}

	model->layers = calloc(model->num_layers, sizeof(qwen3_layer_t));
	if (!model->layers) {
		free(model);
		return NULL;
	}

	/* Open safetensors shards dynamically */
	safetensors_file_t *files[QWEN3_MAX_SHARDS];
	int                 num_files = open_safetensors_shards(model_dir, files, QWEN3_MAX_SHARDS);
	if (num_files == 0) {
		fprintf(stderr, "qwen3_model_load: failed to open safetensors files\n");
		free(model->layers);
		free(model);
		return NULL;
	}

	/* Load embedding weights */
	model->embed_tokens = load_tensor(files, num_files, "model.embed_tokens.weight");
	if (!model->embed_tokens) {
		fprintf(stderr, "qwen3_model_load: failed to load embed_tokens\n");
		goto error;
	}

	/* Load layer weights */
	for (int i = 0; i < model->num_layers; i++) {
		if (load_layer_weights(&model->layers[i], files, num_files, i) != 0) {
			fprintf(stderr, "qwen3_model_load: failed to load layer %d\n", i);
			goto error;
		}
	}

	/* Load final norm */
	model->norm_weight = load_tensor(files, num_files, "model.norm.weight");
	if (!model->norm_weight) {
		fprintf(stderr, "qwen3_model_load: failed to load final norm\n");
		goto error;
	}

	for (int i = 0; i < num_files; i++)
		safetensors_close(files[i]);

	/* Compute RoPE frequencies */
	int max_seq     = QWEN3_MAX_SEQ_LEN;
	int half_dim    = model->head_dim / 2;
	model->rope_cos = malloc(max_seq * half_dim * sizeof(float));
	model->rope_sin = malloc(max_seq * half_dim * sizeof(float));
	compute_rope_freqs(model->rope_cos, model->rope_sin, max_seq, model->head_dim, model->rope_theta);

	/* Allocate working memory */
	qwen3_alloc_work_buffers(model);

	return model;

error:
	for (int i = 0; i < num_files; i++)
		safetensors_close(files[i]);
	qwen3_model_free(model);
	return NULL;
}

/* Load model in mmap mode - keeps safetensors files open and loads layer weights
 * on-demand during forward pass. Reduces peak memory from ~16GB to ~2GB. */
/* Memory-mapped loading: keeps safetensors files open and loads layer weights
 * on demand during forward pass. Only embeddings, final norm, and RoPE tables
 * are resident. Dramatically reduces startup time and peak memory (~2GB vs
 * ~16GB). For GPU: also loads bf16 weight pointers for zero-copy GPU upload
 * directly from the mmap region. */
qwen3_model_t *qwen3_model_load_mmap(const char *model_dir) {
	qwen3_model_t *model = calloc(1, sizeof(qwen3_model_t));
	if (!model)
		return NULL;

	model->use_mmap = 1;

	/* Parse config, fall back to defaults */
	if (parse_qwen3_config(model_dir, model) != 0) {
		qwen3_set_defaults(model);
	}

#ifdef USE_METAL
	/* Enable bf16 GPU acceleration when Metal is available.
	 * Set IRIS_QWEN3_NO_BF16=1 to disable for debugging. */
	model->use_bf16 = (iris_metal_available() && !getenv("IRIS_QWEN3_NO_BF16")) ? 1 : 0;
	if (model->use_bf16) {
		if (iris_verbose)
			fprintf(stderr, "Qwen3: bf16 GPU acceleration enabled\n");
	}
#endif
#ifdef USE_VULKAN
	/* Enable the fully GPU-resident bf16 text-encoder path when Vulkan exposes
	 * the resident Qwen ops. Set IRIS_QWEN3_NO_BF16=1 to disable for debugging. */
	model->use_bf16 = (iris_vk_qwen_available() && !getenv("IRIS_QWEN3_NO_BF16")) ? 1 : 0;
	if (model->use_bf16 && iris_verbose)
		fprintf(stderr, "Qwen3: Vulkan resident bf16 text encoder enabled\n");
#endif
	model->layers = calloc(model->num_layers, sizeof(qwen3_layer_t));
	if (!model->layers) {
		free(model);
		return NULL;
	}

	/* Open safetensors shards dynamically and keep them open */
	model->num_sf_files = open_safetensors_shards(model_dir, model->sf_files, QWEN3_MAX_SHARDS);
	if (model->num_sf_files == 0) {
		fprintf(stderr, "qwen3_model_load_mmap: failed to open safetensors files\n");
		goto error;
	}

#if defined(USE_METAL) || defined(USE_VULKAN)
	/* Detect Q8_0 projection weights from a representative tensor. The whole
	 * checkpoint is homogeneous: all projection weights share one format. */
	{
		char                buf[256];
		const char         *nm = qwen3_lookup_name(model->sf_files[0], "model.layers.0.self_attn.q_proj.weight", buf, sizeof(buf));
		const safetensor_t *t  = safetensors_find(model->sf_files[0], nm);
		model->weights_q8      = (t && safetensor_is_q8_0(t)) ? 1 : 0;
		if (model->weights_q8 && iris_verbose)
			fprintf(stderr, "Qwen3: Q8_0 quantized weights detected\n");
	}
#endif

	/* Load only embeddings - needed for all tokens */
	model->embed_tokens = load_tensor(model->sf_files, model->num_sf_files, "model.embed_tokens.weight");
	if (!model->embed_tokens) {
		fprintf(stderr, "qwen3_model_load_mmap: failed to load embed_tokens\n");
		goto error;
	}

	/* Load final norm (small) */
	model->norm_weight = load_tensor(model->sf_files, model->num_sf_files, "model.norm.weight");
	if (!model->norm_weight) {
		fprintf(stderr, "qwen3_model_load_mmap: failed to load final norm\n");
		goto error;
	}

	/* DON'T load layer weights - they'll be loaded on-demand in forward pass */
	if (iris_verbose)
		fprintf(stderr, "Mmap mode: layer weights will be loaded on-demand\n");

	/* Compute RoPE frequencies */
	int max_seq     = QWEN3_MAX_SEQ_LEN;
	int half_dim    = model->head_dim / 2;
	model->rope_cos = malloc(max_seq * half_dim * sizeof(float));
	model->rope_sin = malloc(max_seq * half_dim * sizeof(float));
	compute_rope_freqs(model->rope_cos, model->rope_sin, max_seq, model->head_dim, model->rope_theta);

	/* Allocate working memory */
	qwen3_alloc_work_buffers(model);

	return model;

error:
	qwen3_model_free(model);
	return NULL;
}

void qwen3_model_free(qwen3_model_t *model) {
	if (!model)
		return;

	free(model->embed_tokens);
	free(model->norm_weight);
	free(model->rope_cos);
	free(model->rope_sin);

	if (model->layers) {
		for (int i = 0; i < model->num_layers; i++) {
			qwen3_layer_t *layer = &model->layers[i];
			free(layer->input_layernorm_weight);
			free(layer->post_attention_layernorm_weight);
			free(layer->attn.q_proj_weight);
			free(layer->attn.k_proj_weight);
			free(layer->attn.v_proj_weight);
			free(layer->attn.o_proj_weight);
			free(layer->attn.q_norm_weight);
			free(layer->attn.k_norm_weight);
			free(layer->mlp.gate_proj_weight);
			free(layer->mlp.up_proj_weight);
			free(layer->mlp.down_proj_weight);
#if defined(USE_METAL) || defined(USE_VULKAN)
			/* Persistent bf16 norm copies (see load_layer_weights_bf16). The
			 * bf16 projection pointers are direct mmap views - do not free. */
			free(layer->input_layernorm_weight_bf16);
			free(layer->post_attention_layernorm_weight_bf16);
			free(layer->attn.q_norm_weight_bf16);
			free(layer->attn.k_norm_weight_bf16);
#endif
		}
		free(model->layers);
	}

	free(model->hidden_state);
	free(model->residual);
	free(model->q_buf);
	free(model->k_buf);
	free(model->v_buf);
	free(model->attn_scores);
	free(model->attn_out);
	free(model->mlp_gate);
	free(model->mlp_up);
	free(model->mlp_out);
	free(model->norm_buf);

	/* Free attention work buffers */
	free(model->attn_q_head);
	free(model->attn_v_head);
	free(model->attn_out_head);

	for (int i = 0; i < 3; i++) {
		free(model->layer_outputs[i]);
	}

	/* Close mmap'd safetensors files if open */
	for (int i = 0; i < model->num_sf_files; i++) {
		if (model->sf_files[i])
			safetensors_close(model->sf_files[i]);
	}

	free(model);
}

/* ========================================================================
 * Combined Encoder API
 * ======================================================================== */

qwen3_encoder_t *qwen3_encoder_load(const char *model_dir, int use_mmap) {
	qwen3_encoder_t *enc = calloc(1, sizeof(qwen3_encoder_t));
	if (!enc)
		return NULL;

	/* Load tokenizer */
	char tok_path[512];
	snprintf(tok_path, sizeof(tok_path), "%s/tokenizer.json", model_dir);
	enc->tokenizer = qwen3_tokenizer_load(tok_path);
	if (!enc->tokenizer) {
		fprintf(stderr, "qwen3_encoder_load: failed to load tokenizer\n");
		free(enc);
		return NULL;
	}

	/* Load model - use mmap mode if requested (saves ~14GB RAM).
	 * The GGUF checkpoint now lives directly in the model root. */
	const char *model_path = model_dir;
	if (use_mmap) {
		enc->model = qwen3_model_load_mmap(model_path);
	}
	else {
		enc->model = qwen3_model_load(model_path);
	}
	if (!enc->model) {
		fprintf(stderr, "qwen3_encoder_load: failed to load model\n");
		qwen3_tokenizer_free(enc->tokenizer);
		free(enc);
		return NULL;
	}

	return enc;
}

void qwen3_encoder_free(qwen3_encoder_t *enc) {
	if (!enc)
		return;
	qwen3_tokenizer_free(enc->tokenizer);
	qwen3_model_free(enc->model);
	free(enc);
}

/* Main text encoding API. Tokenizes the prompt using the Qwen3 chat template
 * (with <think> tags), pads to max sequence length, runs the forward pass, and
 * returns extracted embeddings. Also returns the number of real (non-padding)
 * tokens via out_num_tokens. */
float *qwen3_encode_text_ex(qwen3_encoder_t *enc, const char *prompt, int *out_num_tokens) {
	if (!enc || !enc->tokenizer || !enc->model || !prompt)
		return NULL;

	/* Tokenize with chat template */
	int  num_tokens;
	int *tokens = qwen3_tokenize_chat(enc->tokenizer, prompt, &num_tokens, QWEN3_MAX_SEQ_LEN);
	if (!tokens)
		return NULL;

	/* Pad to max length */
	int *attention_mask = malloc(QWEN3_MAX_SEQ_LEN * sizeof(int));
	int *padded_tokens  = qwen3_pad_tokens(tokens, num_tokens, QWEN3_MAX_SEQ_LEN, attention_mask);
	free(tokens);

	if (!padded_tokens) {
		free(attention_mask);
		return NULL;
	}

	if (out_num_tokens)
		*out_num_tokens = num_tokens;

	/* Forward pass */
	float *embeddings = qwen3_forward(enc->model, padded_tokens, attention_mask, QWEN3_MAX_SEQ_LEN);

	free(padded_tokens);
	free(attention_mask);

	return embeddings;
}

float *qwen3_encode_text(qwen3_encoder_t *enc, const char *prompt) {
	return qwen3_encode_text_ex(enc, prompt, NULL);
}