diff --git a/.gitignore b/.gitignore index ef4df4c..21b62b0 100644 --- a/.gitignore +++ b/.gitignore @@ -62,6 +62,9 @@ data/ # Tokenizer data files (small, needed for inference) !tokenizer/data/ +tokenizer/data/orion_*/ +tokenizer/data/*_tok/ +tokenizer/hf_*/ # Experiments (binaries only, source tracked) experiments/*.o diff --git a/Makefile b/Makefile index ac8e82d..c34c2e9 100644 --- a/Makefile +++ b/Makefile @@ -29,11 +29,14 @@ INFERENCE_SRC = \ kernels/inference/prefill_ane.m \ kernels/inference/decode_ane.m \ kernels/inference/decode_cpu.m \ - kernels/inference/kv_cache.m + kernels/inference/kv_cache.m \ + kernels/inference/qwen_cpu_ops.m TRAINING_SRC = \ kernels/training/stories_train.m \ kernels/training/stories_cpu_ops.m \ + kernels/training/qwen_lora_cpu_ops.m \ + kernels/training/qwen_lora_train.m \ kernels/training/data_loader.m MODEL_SRC = model/weight_loader.m @@ -65,6 +68,7 @@ COMPILER_C_SRC = \ compiler/frontends/gpt2_prefill.c \ compiler/frontends/gpt2_decode.c \ compiler/frontends/gpt2_final.c \ + compiler/frontends/qwen35_prefill.c \ compiler/frontends/classifier_softmax.c \ compiler/frontends/stories_train.c \ compiler/frontends/lora.c diff --git a/README_ORION_Q.md b/README_ORION_Q.md new file mode 100644 index 0000000..1721f5f --- /dev/null +++ b/README_ORION_Q.md @@ -0,0 +1,131 @@ +# Orion-Q + +Orion-Q is a Qwen-focused porting and diagnostics subset built on top of Orion. + +It is not a separate engine. It is a curated extension of Orion that adds: + +- Qwen model configs and blob conversion +- Qwen CPU inference path +- Qwen ANE or hybrid inference path +- Qwen LoRA training primitives +- Qwen diagnostics, including smoke, probe, parity, and diff tests + +The boundary for this subset is defined in: + +- `docs/orion_q/ADR-008-orion-q-boundary.md` + +## What Orion-Q Is + +Orion-Q is the part of the local Orion worktree that makes Qwen-family models runnable and verifiable inside Orion. + +In practical terms, Orion-Q includes: + +- Qwen frontend and model registration +- Qwen weight loading and export path +- Qwen-specific CPU and ANE execution support +- Qwen LoRA training path +- Qwen diagnostics and validation tests + +## What Orion-Q Is Not + +Orion-Q does not include: + +- Silver accelerator work +- user-specific training tracks +- CRPG or other domain assets +- reports, logs, or generated tokenizer experiment outputs +- exported checkpoints or model weights + +Those belong to downstream tracks or local runtime artifacts, not to Orion-Q itself. + +## Current Status + +Within the currently defined Orion-Q scope: + +- Qwen porting core: complete +- binary judge diagnostics: close-out achieved +- target hybrid parity smoke scope: close-out achieved +- ANE training preparation line: complete + +Supporting documents: + +- `docs/orion_q/ORION_Q_PORT_CLOSEOUT.md` +- `docs/orion_q/ORION_Q_HYBRID_PARITY_CLOSEOUT.md` +- `docs/orion_q/ORION_Q_ANE_PREP_CLOSEOUT.md` + +## Included Code Areas + +The shared Orion-Q subset is expected to cover these groups: + +- shared Orion core changes required by Qwen support +- `compiler/frontends/qwen35_*` +- `kernels/inference/qwen_*` +- `kernels/training/qwen_lora_*` +- `model/configs/qwen35_*` +- `model/convert/hf_to_blobs_qwen35.py` +- `tests/test_qwen35_*` +- `tests/test_qwen35_9b_*` + +## Validation Philosophy + +Orion-Q treats diagnostics as part of the product surface, not as throwaway experiments. + +That means the following are first-class parts of the subset: + +- smoke tests +- bridge-stage diffs +- layer diffs +- parity checks +- ANE training probes + +## Recommended Share Mode + +The recommended way to share Orion-Q is: + +1. As an Orion-based fork or draft PR branch +2. With generated artifacts excluded +3. With a narrow, explicit scope + +Suggested framing: + +`Orion-Q: a Qwen-focused porting and diagnostics subset built on top of Orion` + +## Excluded Artifacts + +Do not publish these as part of Orion-Q: + +- `tokenizer/data/orion_*` +- `tokenizer/data/*_tok` +- exported blobs and checkpoints +- local reports and logs +- user workflow assets + +## Local Share Bundle + +This workspace can generate a clean Orion-Q share bundle with: + +```bash +python3 scripts/prepare_orion_q_share.py --clean +``` + +Default output: + +```text +build/orion_q_share/Orion-Q +``` + +The bundle is driven by: + +- `scripts/orion_q_share_manifest.txt` + +## Relationship to Upstream Orion + +Upstream Orion remains the execution core. + +Orion-Q should be communicated as: + +- Orion core +- plus Qwen-specific porting +- plus Qwen-specific diagnostics + +It should not be presented as a replacement brand or a disconnected new project. diff --git a/apps/cli/commands/infer.m b/apps/cli/commands/infer.m index 80d9d99..aff1439 100644 --- a/apps/cli/commands/infer.m +++ b/apps/cli/commands/infer.m @@ -121,6 +121,13 @@ int orion_cmd_infer(int argc, const char* argv[]) { double t_load = time_ms() - t0; fprintf(stderr, "Weights loaded in %.1f ms\n", t_load); + OrionModelConfig cfg = kGPT2_124M; + cfg.n_layer = w->n_layer; + cfg.d_model = w->d_model; + cfg.hidden_dim = w->d_ff; + cfg.vocab = w->vocab; + cfg.max_seq = w->max_seq; + // Tokenize prompt int prompt_tokens[1024]; int prompt_len = orion_gpt2_encode(tok, prompt, prompt_tokens, 1024); @@ -134,7 +141,7 @@ int orion_cmd_infer(int argc, const char* argv[]) { // Allocate float* logits = (float*)malloc(w->vocab * sizeof(float)); - OrionKVCache* kv = orion_kv_cache_create(&kGPT2_124M); + OrionKVCache* kv = orion_kv_cache_create(&cfg); int gen_count = 0; // Start profiler @@ -148,12 +155,12 @@ int orion_cmd_infer(int argc, const char* argv[]) { bool prefill_ok; if (use_ane) { prefill_ok = orion_ane_prefill(w, prompt_tokens, prompt_len, - &kGPT2_124M, weights_path, kv, logits); + &cfg, weights_path, kv, logits); if (!prefill_ok) { fprintf(stderr, "Warning: ANE prefill failed, falling back to CPU\n"); // Reset KV cache orion_kv_cache_free(kv); - kv = orion_kv_cache_create(&kGPT2_124M); + kv = orion_kv_cache_create(&cfg); orion_gpt2_prefill_kv(w, prompt_tokens, prompt_len, kv, logits); prefill_ok = true; } diff --git a/compiler/builder.c b/compiler/builder.c index b2ca9a5..0e76fad 100644 --- a/compiler/builder.c +++ b/compiler/builder.c @@ -4,6 +4,7 @@ #include #include #include +#include // Helper: create a node with basic fields static OrionNode make_node(OrionOp op, const char* name, OrionDtype dtype, int shape[4]) { @@ -17,6 +18,11 @@ static OrionNode make_node(OrionOp op, const char* name, OrionDtype dtype, int s return n; } +static bool orion_use_fp32_rrms_powchain(void) { + const char *mode = getenv("ORION_RMSNORM_RRMS_MODE"); + return mode && (strcmp(mode, "nr1") == 0 || strcmp(mode, "fp32") == 0); +} + int orion_gb_input(OrionGraph* g, const char* name, OrionDtype dtype, int shape[4]) { OrionNode n = make_node(ORION_OP_INPUT, name, dtype, shape); int idx = orion_graph_add_node(g, &n); @@ -335,8 +341,24 @@ int orion_gb_gelu(OrionGraph* g, int input, const char* prefix, int dim __attrib int orion_gb_silu(OrionGraph* g, int input, const char* prefix, int dim __attribute__((unused)), int seq __attribute__((unused))) { char buf[ORION_MAX_NAME]; + // Use sigmoid(x) = 0.5 * (tanh(0.5 * x) + 1) to avoid the ANE builtin sigmoid drift. + snprintf(buf, sizeof(buf), "%s_half", prefix); + int half = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, 0.5f); + + snprintf(buf, sizeof(buf), "%s_hx", prefix); + int hx = orion_gb_mul(g, input, half, buf); + + snprintf(buf, sizeof(buf), "%s_th", prefix); + int th = orion_gb_tanh(g, hx, buf); + + snprintf(buf, sizeof(buf), "%s_one", prefix); + int one = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, 1.0f); + + snprintf(buf, sizeof(buf), "%s_onep", prefix); + int onep = orion_gb_add(g, th, one, buf); + snprintf(buf, sizeof(buf), "%s_sig", prefix); - int sig = orion_gb_sigmoid(g, input, buf); + int sig = orion_gb_mul(g, onep, half, buf); snprintf(buf, sizeof(buf), "%s_out", prefix); int out = orion_gb_mul(g, input, sig, buf); @@ -366,6 +388,7 @@ int orion_gb_rmsnorm(OrionGraph* g, int input, int weight, float eps, // / dim float inv_dim = 1.0f / (float)dim; snprintf(buf, sizeof(buf), "%s_invd", prefix); + snprintf(buf, sizeof(buf), "%s_invd", prefix); int invd = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, inv_dim); snprintf(buf, sizeof(buf), "%s_ms", prefix); int ms = orion_gb_mul(g, ss, invd, buf); @@ -382,9 +405,29 @@ int orion_gb_rmsnorm(OrionGraph* g, int input, int weight, float eps, snprintf(buf, sizeof(buf), "%s_rrms", prefix); int rrms = orion_gb_pow(g, mse, nhalf, buf); - // x * rrms - snprintf(buf, sizeof(buf), "%s_xr", prefix); - int xr = orion_gb_mul(g, input, rrms, buf); + int xr = -1; + if (orion_use_fp32_rrms_powchain()) { + snprintf(buf, sizeof(buf), "%s_rrms_sq", prefix); + int rrms_sq = orion_gb_mul(g, rrms, rrms, buf); + snprintf(buf, sizeof(buf), "%s_nr_term", prefix); + int nr_term = orion_gb_mul(g, mse, rrms_sq, buf); + snprintf(buf, sizeof(buf), "%s_half_nr", prefix); + int half_nr = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, 0.5f); + snprintf(buf, sizeof(buf), "%s_nr_half", prefix); + int nr_half = orion_gb_mul(g, nr_term, half_nr, buf); + snprintf(buf, sizeof(buf), "%s_threehalves", prefix); + int threehalves = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, 1.5f); + snprintf(buf, sizeof(buf), "%s_nr_corr", prefix); + int nr_corr = orion_gb_sub(g, threehalves, nr_half, buf); + snprintf(buf, sizeof(buf), "%s_rrms_refined", prefix); + int rrms_refined = orion_gb_mul(g, rrms, nr_corr, buf); + snprintf(buf, sizeof(buf), "%s_xr", prefix); + xr = orion_gb_mul(g, input, rrms_refined, buf); + } else { + // x * rrms + snprintf(buf, sizeof(buf), "%s_xr", prefix); + xr = orion_gb_mul(g, input, rrms, buf); + } // weight * normalized snprintf(buf, sizeof(buf), "%s_out", prefix); diff --git a/compiler/frontends/qwen35_prefill.c b/compiler/frontends/qwen35_prefill.c new file mode 100644 index 0000000..e2c23f2 --- /dev/null +++ b/compiler/frontends/qwen35_prefill.c @@ -0,0 +1,180 @@ +// compiler/frontends/qwen35_prefill.c — Qwen3.5 ANE prefill frontend + +#include "qwen35_prefill.h" +#include "../builder.h" +#include "../patterns.h" +#include +#include +#include + +static int qwen35_kv_dim(const OrionModelConfig* cfg) { + int n_kv = cfg->n_kv_head > 0 ? cfg->n_kv_head : cfg->n_head; + return n_kv * cfg->head_dim; +} + +static int qwen35_use_matmul_v_proj(void) { + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "matmul") == 0; +} + +static int qwen35_input_rmsnorm(OrionGraph* g, int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_pattern_cast_to_fp16(g, x, "x16", d, s); + + char path[256]; + int ln_shape[4] = {1, d, 1, 1}; + snprintf(path, sizeof(path), "@model_path/layer%d/input_layernorm.bin", layer); + int rms_w = orion_gb_const_weight(g, "input_ln_w", ORION_DTYPE_FP16, ln_shape, path, 64); + return orion_gb_rmsnorm(g, x16, rms_w, 1e-6f, "input_rms", d, s); +} + +static int qwen35_normed_input(OrionGraph* g, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + return orion_pattern_cast_to_fp16(g, x, "x16", d, s); +} + +static int qwen35_linear_matmul(OrionGraph* g, + int input, + const char* prefix, + int in_dim, + int out_dim, + int seq, + const char* weight_path) { + char buf[256]; + int perm_shape[4] = {4, 1, 1, 1}; + int perm_vals[4] = {0, 2, 3, 1}; + int perm_back_vals[4] = {0, 3, 1, 2}; + + snprintf(buf, sizeof(buf), "%s_pm", prefix); + int pm = orion_gb_const_int32(g, buf, perm_shape, perm_vals, 4); + snprintf(buf, sizeof(buf), "%s_pmb", prefix); + int pm_back = orion_gb_const_int32(g, buf, perm_shape, perm_back_vals, 4); + + int xt_shape[4] = {1, 1, seq, in_dim}; + snprintf(buf, sizeof(buf), "%s_xt", prefix); + int xt = orion_gb_transpose(g, input, pm, buf, perm_vals, xt_shape); + + int w_shape[4] = {1, 1, out_dim, in_dim}; + snprintf(buf, sizeof(buf), "%s_W", prefix); + int w = orion_gb_const_weight(g, buf, ORION_DTYPE_FP16, w_shape, weight_path, 64); + + int mm_shape[4] = {1, 1, seq, out_dim}; + snprintf(buf, sizeof(buf), "%s_mm", prefix); + int mm = orion_gb_matmul(g, xt, w, false, true, buf, mm_shape); + + int out_shape[4] = {1, out_dim, 1, seq}; + snprintf(buf, sizeof(buf), "%s_out", prefix); + return orion_gb_transpose(g, mm, pm_back, buf, perm_back_vals, out_shape); +} + +static int qwen35_post_attn_rmsnorm(OrionGraph* g, int input, int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int s = bucket; + char path[256]; + int ln_shape[4] = {1, d, 1, 1}; + snprintf(path, sizeof(path), "@model_path/layer%d/post_attention_layernorm.bin", layer); + int rms_w = orion_gb_const_weight(g, "post_attn_ln_w", ORION_DTYPE_FP16, ln_shape, path, 64); + return orion_gb_rmsnorm(g, input, rms_w, 1e-6f, "post_attn_rms", d, s); +} + +OrionGraph* orion_frontend_qwen35_prefill_q_proj(int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int s = bucket; + OrionGraph* g = orion_graph_create(); + int rms = qwen35_input_rmsnorm(g, layer, bucket, cfg); + + char path[256]; + snprintf(path, sizeof(path), "@model_path/layer%d/self_attn_q_proj.bin", layer); + int q_proj = orion_gb_linear(g, rms, "q_proj", d, d * 2, s, path, NULL); + int q_proj32 = orion_pattern_cast_to_fp32(g, q_proj, "q_proj32", d * 2, s); + orion_gb_output(g, q_proj32, "q_proj"); + return g; +} + +OrionGraph* orion_frontend_qwen35_prefill_q_proj_linear_only(int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int s = bucket; + OrionGraph* g = orion_graph_create(); + int x16 = qwen35_normed_input(g, bucket, cfg); + + char path[256]; + snprintf(path, sizeof(path), "@model_path/layer%d/self_attn_q_proj.bin", layer); + int q_proj = orion_gb_linear(g, x16, "q_proj", d, d * 2, s, path, NULL); + int q_proj32 = orion_pattern_cast_to_fp32(g, q_proj, "q_proj32", d * 2, s); + orion_gb_output(g, q_proj32, "q_proj"); + return g; +} + +OrionGraph* orion_frontend_qwen35_prefill_kv_proj(int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int kv = qwen35_kv_dim(cfg); + int s = bucket; + OrionGraph* g = orion_graph_create(); + int rms = qwen35_input_rmsnorm(g, layer, bucket, cfg); + + char k_path[256], v_path[256]; + snprintf(k_path, sizeof(k_path), "@model_path/layer%d/self_attn_k_proj.bin", layer); + snprintf(v_path, sizeof(v_path), "@model_path/layer%d/self_attn_v_proj.bin", layer); + + int k_proj = orion_gb_linear(g, rms, "k_proj", d, kv, s, k_path, NULL); + int v_proj = qwen35_use_matmul_v_proj() + ? qwen35_linear_matmul(g, rms, "v_proj", d, kv, s, v_path) + : orion_gb_linear(g, rms, "v_proj", d, kv, s, v_path, NULL); + int k_proj32 = orion_pattern_cast_to_fp32(g, k_proj, "k_proj32", kv, s); + int v_proj32 = orion_pattern_cast_to_fp32(g, v_proj, "v_proj32", kv, s); + + orion_gb_output(g, k_proj32, "k_proj"); + orion_gb_output(g, v_proj32, "v_proj"); + return g; +} + +OrionGraph* orion_frontend_qwen35_prefill_kv_proj_linear_only(int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int kv = qwen35_kv_dim(cfg); + int s = bucket; + OrionGraph* g = orion_graph_create(); + int x16 = qwen35_normed_input(g, bucket, cfg); + + char k_path[256], v_path[256]; + snprintf(k_path, sizeof(k_path), "@model_path/layer%d/self_attn_k_proj.bin", layer); + snprintf(v_path, sizeof(v_path), "@model_path/layer%d/self_attn_v_proj.bin", layer); + + int k_proj = orion_gb_linear(g, x16, "k_proj", d, kv, s, k_path, NULL); + int v_proj = qwen35_use_matmul_v_proj() + ? qwen35_linear_matmul(g, x16, "v_proj", d, kv, s, v_path) + : orion_gb_linear(g, x16, "v_proj", d, kv, s, v_path, NULL); + int k_proj32 = orion_pattern_cast_to_fp32(g, k_proj, "k_proj32", kv, s); + int v_proj32 = orion_pattern_cast_to_fp32(g, v_proj, "v_proj32", kv, s); + + orion_gb_output(g, k_proj32, "k_proj"); + orion_gb_output(g, v_proj32, "v_proj"); + return g; +} + +OrionGraph* orion_frontend_qwen35_prefill_ffn(int layer, int bucket, const OrionModelConfig* cfg) { + int d = cfg->d_model; + int h = cfg->hidden_dim; + int s = bucket; + OrionGraph* g = orion_graph_create(); + int in_shape[4] = {1, d, 1, s}; + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_pattern_cast_to_fp16(g, x, "x16", d, s); + int rms = qwen35_post_attn_rmsnorm(g, x16, layer, bucket, cfg); + + char gate_w[256], up_w[256], down_w[256]; + snprintf(gate_w, sizeof(gate_w), "@model_path/layer%d/mlp_gate_proj.bin", layer); + snprintf(up_w, sizeof(up_w), "@model_path/layer%d/mlp_up_proj.bin", layer); + snprintf(down_w, sizeof(down_w), "@model_path/layer%d/mlp_down_proj.bin", layer); + + int ffn = orion_pattern_swiglu_ffn(g, rms, "ffn", d, h, s, gate_w, up_w, down_w); + int resid = orion_pattern_residual(g, x16, ffn, "ffn_resid"); + int hidden = orion_pattern_cast_to_fp32(g, resid, "hidden", d, s); + orion_gb_output(g, hidden, "hidden"); + return g; +} diff --git a/compiler/frontends/qwen35_prefill.h b/compiler/frontends/qwen35_prefill.h new file mode 100644 index 0000000..057ecf9 --- /dev/null +++ b/compiler/frontends/qwen35_prefill.h @@ -0,0 +1,35 @@ +// compiler/frontends/qwen35_prefill.h — Qwen3.5 ANE prefill frontend +#ifndef ORION_FRONTEND_QWEN35_PREFILL_H +#define ORION_FRONTEND_QWEN35_PREFILL_H + +#include "../graph.h" +#include "../model_config.h" + +// Build a single-output Q projection graph for one full-attention layer. +// Input: hidden [1, d_model, 1, seq] fp32 +// Output: q_proj [1, 2*d_model, 1, seq] fp32 +OrionGraph* orion_frontend_qwen35_prefill_q_proj(int layer, int bucket, const OrionModelConfig* cfg); + +// Build a single-output Q projection graph that assumes the input is already RMSNorm'd. +// Input: normed hidden [1, d_model, 1, seq] fp32 +// Output: q_proj [1, 2*d_model, 1, seq] fp32 +OrionGraph* orion_frontend_qwen35_prefill_q_proj_linear_only(int layer, int bucket, const OrionModelConfig* cfg); + +// Build a K/V projection graph for one full-attention layer. +// Input: hidden [1, d_model, 1, seq] fp32 +// Outputs: k_proj [1, n_kv_head*head_dim, 1, seq] fp32 +// v_proj [1, n_kv_head*head_dim, 1, seq] fp32 +OrionGraph* orion_frontend_qwen35_prefill_kv_proj(int layer, int bucket, const OrionModelConfig* cfg); + +// Build a K/V projection graph that assumes the input is already RMSNorm'd. +// Input: normed hidden [1, d_model, 1, seq] fp32 +// Outputs: k_proj [1, n_kv_head*head_dim, 1, seq] fp32 +// v_proj [1, n_kv_head*head_dim, 1, seq] fp32 +OrionGraph* orion_frontend_qwen35_prefill_kv_proj_linear_only(int layer, int bucket, const OrionModelConfig* cfg); + +// Build a post-attention FFN graph for one Qwen3.5 layer. +// Input: hidden [1, d_model, 1, seq] fp32 +// Output: hidden [1, d_model, 1, seq] fp32 +OrionGraph* orion_frontend_qwen35_prefill_ffn(int layer, int bucket, const OrionModelConfig* cfg); + +#endif // ORION_FRONTEND_QWEN35_PREFILL_H diff --git a/compiler/model_config.h b/compiler/model_config.h index 3a6b54c..3534907 100644 --- a/compiler/model_config.h +++ b/compiler/model_config.h @@ -5,10 +5,12 @@ // Mirror of OrionModelConfig from core/ane_runtime.h // Keep in sync manually. Both are POD structs with identical layout. -#ifndef ORION_ANE_RUNTIME_H +#ifndef ORION_MODEL_CONFIG_DEFINED +#define ORION_MODEL_CONFIG_DEFINED typedef struct { int n_layer; int n_head; + int n_kv_head; int d_model; int head_dim; int hidden_dim; diff --git a/core/ane_runtime.h b/core/ane_runtime.h index 05e5294..27741eb 100644 --- a/core/ane_runtime.h +++ b/core/ane_runtime.h @@ -5,6 +5,8 @@ #import /// Model configuration shared across all Orion components. +#ifndef ORION_MODEL_CONFIG_DEFINED +#define ORION_MODEL_CONFIG_DEFINED typedef struct { int n_layer; int n_head; @@ -13,7 +15,9 @@ typedef struct { int hidden_dim; int vocab; int max_seq; + int n_kv_head; } OrionModelConfig; +#endif /// Opaque handle to a compiled ANE program. typedef struct OrionProgram OrionProgram; @@ -55,6 +59,11 @@ void orion_release_program(OrionProgram* prog); /// Useful for tracking approach to the ~119 compile limit. int orion_compile_count(void); +/// Return the temp directory used by a compiled ANE program, or nil. +/// The returned NSString is owned by the program and remains valid until +/// orion_release_program() is called for that program. +NSString* orion_program_tmp_dir(OrionProgram* prog); + /// Create a new ANE program with patched weights, reusing compiled artifacts /// from a donor program. Skips compilation entirely — only loads. /// @@ -91,4 +100,27 @@ bool orion_program_reload_weights( NSDictionary* weight_dict ); +/// Export a loaded program's compiled runtime artifacts to a persistent directory. +/// The destination is replaced atomically enough for single-writer Silver cache usage. +/// @param prog Loaded OrionProgram +/// @param artifact_dir Persistent directory to populate +/// @return true on success. +bool orion_program_export_artifacts( + OrionProgram* prog, + const char* artifact_dir +); + +/// Load an ANE program from previously exported runtime artifacts without compiling. +/// @param mil_text Same MIL text used for the original compile +/// @param weight_dict Same weight dict key structure used for the original compile +/// @param artifact_dir Directory containing exported runtime artifacts +/// @param program_tag Tag for debugging (may be NULL) +/// @return Loaded program handle, or NULL on failure. +OrionProgram* orion_program_load_artifacts( + const char* mil_text, + NSDictionary* weight_dict, + const char* artifact_dir, + const char* program_tag +); + #endif // ORION_ANE_RUNTIME_H diff --git a/core/ane_runtime.m b/core/ane_runtime.m index 98208c8..3aa7151 100644 --- a/core/ane_runtime.m +++ b/core/ane_runtime.m @@ -2,6 +2,7 @@ #import #import #import +#import // T015: orion_compile_mil — compile MIL text to ANE program // T016: orion_eval — evaluate a compiled program @@ -218,6 +219,11 @@ int orion_compile_count(void) { return g_compile_count; } +NSString* orion_program_tmp_dir(OrionProgram* prog) { + if (!prog || !prog->tmpDir) return nil; + return (__bridge NSString *)prog->tmpDir; +} + #pragma mark - T151: Delta weight patching OrionProgram* orion_program_patch_weights( @@ -384,3 +390,123 @@ bool orion_program_reload_weights( return true; } } + +static bool _orion_copy_item_if_exists(NSFileManager *fm, NSString *src, NSString *dst) { + if (![fm fileExistsAtPath:src]) return true; + [fm removeItemAtPath:dst error:nil]; + NSString *parent = [dst stringByDeletingLastPathComponent]; + [fm createDirectoryAtPath:parent withIntermediateDirectories:YES attributes:nil error:nil]; + NSError *err = nil; + return [fm copyItemAtPath:src toPath:dst error:&err]; +} + +bool orion_program_export_artifacts( + OrionProgram* prog, + const char* artifact_dir +) { + if (!prog || !prog->tmpDir || !artifact_dir) return false; + + @autoreleasepool { + NSFileManager *fm = [NSFileManager defaultManager]; + NSString *srcDir = (__bridge NSString *)prog->tmpDir; + NSString *dstDir = [NSString stringWithUTF8String:artifact_dir]; + NSString *stagingDir = [dstDir stringByAppendingFormat:@".tmp.%d", getpid()]; + + [fm removeItemAtPath:stagingDir error:nil]; + [fm createDirectoryAtPath:stagingDir withIntermediateDirectories:YES attributes:nil error:nil]; + + bool ok = true; + ok = ok && _orion_copy_item_if_exists(fm, + [srcDir stringByAppendingPathComponent:@"data"], + [stagingDir stringByAppendingPathComponent:@"data"]); + ok = ok && _orion_copy_item_if_exists(fm, + [srcDir stringByAppendingPathComponent:@"net.plist"], + [stagingDir stringByAppendingPathComponent:@"net.plist"]); + ok = ok && _orion_copy_item_if_exists(fm, + [srcDir stringByAppendingPathComponent:@"model.mil"], + [stagingDir stringByAppendingPathComponent:@"model.mil"]); + ok = ok && _orion_copy_item_if_exists(fm, + [srcDir stringByAppendingPathComponent:@"weights"], + [stagingDir stringByAppendingPathComponent:@"weights"]); + if (!ok) { + [fm removeItemAtPath:stagingDir error:nil]; + return false; + } + + [fm removeItemAtPath:dstDir error:nil]; + NSError *moveErr = nil; + if (![fm moveItemAtPath:stagingDir toPath:dstDir error:&moveErr]) { + [fm removeItemAtPath:stagingDir error:nil]; + return false; + } + return true; + } +} + +OrionProgram* orion_program_load_artifacts( + const char* mil_text, + NSDictionary* weight_dict, + const char* artifact_dir, + const char* program_tag +) { + if (!g_init || !mil_text || !artifact_dir) return NULL; + + @autoreleasepool { + NSString *artifactDir = [NSString stringWithUTF8String:artifact_dir]; + NSFileManager *fm = [NSFileManager defaultManager]; + NSString *artifactData = [artifactDir stringByAppendingPathComponent:@"data"]; + NSString *artifactPlist = [artifactDir stringByAppendingPathComponent:@"net.plist"]; + if (![fm fileExistsAtPath:artifactData] || ![fm fileExistsAtPath:artifactPlist]) { + return NULL; + } + + NSData *milData = [NSData dataWithBytes:mil_text length:strlen(mil_text)]; + NSDictionary *wdict = weight_dict ?: @{}; + + id desc = ((id(*)(Class,SEL,id,id,id))objc_msgSend)( + g_Desc, @selector(modelWithMILText:weights:optionsPlist:), + milData, wdict, nil); + if (!desc) return NULL; + + id model = ((id(*)(Class,SEL,id))objc_msgSend)( + g_IMM, @selector(inMemoryModelWithDescriptor:), desc); + if (!model) return NULL; + + id hexId = ((id(*)(id,SEL))objc_msgSend)(model, @selector(hexStringIdentifier)); + NSString *modelDir = [NSTemporaryDirectory() stringByAppendingPathComponent:hexId]; + + [fm removeItemAtPath:modelDir error:nil]; + [fm createDirectoryAtPath:modelDir withIntermediateDirectories:YES attributes:nil error:nil]; + + bool ok = true; + ok = ok && _orion_copy_item_if_exists(fm, artifactData, [modelDir stringByAppendingPathComponent:@"data"]); + ok = ok && _orion_copy_item_if_exists(fm, artifactPlist, [modelDir stringByAppendingPathComponent:@"net.plist"]); + ok = ok && _orion_copy_item_if_exists(fm, + [artifactDir stringByAppendingPathComponent:@"model.mil"], + [modelDir stringByAppendingPathComponent:@"model.mil"]); + ok = ok && _orion_copy_item_if_exists(fm, + [artifactDir stringByAppendingPathComponent:@"weights"], + [modelDir stringByAppendingPathComponent:@"weights"]); + if (!ok) { + [fm removeItemAtPath:modelDir error:nil]; + return NULL; + } + + NSError *e = nil; + BOOL loadOk = ((BOOL(*)(id,SEL,unsigned int,id,NSError**))objc_msgSend)( + model, @selector(loadWithQoS:options:error:), 21, @{}, &e); + if (!loadOk) { + [fm removeItemAtPath:modelDir error:nil]; + return NULL; + } + + OrionProgram *prog = (OrionProgram *)calloc(1, sizeof(OrionProgram)); + prog->model = (void *)CFBridgingRetain(model); + prog->tmpDir = (void *)CFBridgingRetain(modelDir); + prog->loaded = true; + if (program_tag) { + strlcpy(prog->tag, program_tag, sizeof(prog->tag)); + } + return prog; + } +} diff --git a/core/mil_builder.m b/core/mil_builder.m index c517302..faf2bd0 100644 --- a/core/mil_builder.m +++ b/core/mil_builder.m @@ -1,5 +1,7 @@ #import "mil_builder.h" #import +#include +#include // T019: orion_mil_linear (conv-based) // T020: orion_mil_layernorm, orion_mil_rmsnorm @@ -17,6 +19,11 @@ " tensor %s_dl = const()[name=string(\"%s_dl\"), val=tensor([1,1])];\n" \ " int32 %s_gr = const()[name=string(\"%s_gr\"), val=int32(1)];\n" +static BOOL orion_mil_use_fp32_rrms_powchain(void) { + const char *mode = getenv("ORION_RMSNORM_RRMS_MODE"); + return mode && (strcmp(mode, "nr1") == 0 || strcmp(mode, "fp32") == 0); +} + #pragma mark - T023: Header + Program Wrapper NSString* orion_mil_header(void) { @@ -176,21 +183,47 @@ [m appendFormat:@" bool %@_kd = const()[name=string(\"%@_kd\"), val=bool(true)];\n", p, p]; [m appendFormat:@" tensor %@_ss = reduce_sum(x=%@_sq, axes=%@_ax, keep_dims=%@_kd)[name=string(\"%@_ss\")];\n", seq, p, p, p, p, p]; - [m appendFormat:@" fp16 %@_invd = const()[name=string(\"%@_invd\"), val=fp16(%f)];\n", p, p, inv_dim]; - [m appendFormat:@" tensor %@_ms = mul(x=%@_ss, y=%@_invd)[name=string(\"%@_ms\")];\n", - seq, p, p, p, p]; - - // rsqrt(ms + eps) - [m appendFormat:@" fp16 %@_eps = const()[name=string(\"%@_eps\"), val=fp16(%f)];\n", p, p, eps]; - [m appendFormat:@" tensor %@_mse = add(x=%@_ms, y=%@_eps)[name=string(\"%@_mse\")];\n", - seq, p, p, p, p]; - [m appendFormat:@" fp16 %@_nhalf = const()[name=string(\"%@_nhalf\"), val=fp16(-0.5)];\n", p, p]; - [m appendFormat:@" tensor %@_rrms = pow(x=%@_mse, y=%@_nhalf)[name=string(\"%@_rrms\")];\n", - seq, p, p, p, p]; - - // x * rrms - [m appendFormat:@" tensor %@_xr = mul(x=%@, y=%@_rrms)[name=string(\"%@_xr\")];\n", - dim, seq, p, inp, p, p]; + if (orion_mil_use_fp32_rrms_powchain()) { + [m appendFormat:@" fp16 %@_invd = const()[name=string(\"%@_invd\"), val=fp16(%f)];\n", p, p, inv_dim]; + [m appendFormat:@" tensor %@_ms = mul(x=%@_ss, y=%@_invd)[name=string(\"%@_ms\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" fp16 %@_eps = const()[name=string(\"%@_eps\"), val=fp16(%f)];\n", p, p, eps]; + [m appendFormat:@" tensor %@_mse = add(x=%@_ms, y=%@_eps)[name=string(\"%@_mse\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" fp16 %@_nhalf = const()[name=string(\"%@_nhalf\"), val=fp16(-0.5)];\n", p, p]; + [m appendFormat:@" tensor %@_rrms = pow(x=%@_mse, y=%@_nhalf)[name=string(\"%@_rrms\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" tensor %@_rrms_sq = mul(x=%@_rrms, y=%@_rrms)[name=string(\"%@_rrms_sq\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" tensor %@_nr_term = mul(x=%@_mse, y=%@_rrms_sq)[name=string(\"%@_nr_term\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" fp16 %@_half_nr = const()[name=string(\"%@_half_nr\"), val=fp16(0.5)];\n", p, p]; + [m appendFormat:@" tensor %@_nr_half = mul(x=%@_nr_term, y=%@_half_nr)[name=string(\"%@_nr_half\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" fp16 %@_threehalves = const()[name=string(\"%@_threehalves\"), val=fp16(1.5)];\n", p, p]; + [m appendFormat:@" tensor %@_nr_corr = sub(x=%@_threehalves, y=%@_nr_half)[name=string(\"%@_nr_corr\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" tensor %@_rrms_refined = mul(x=%@_rrms, y=%@_nr_corr)[name=string(\"%@_rrms_refined\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" tensor %@_xr = mul(x=%@, y=%@_rrms_refined)[name=string(\"%@_xr\")];\n", + dim, seq, p, inp, p, p]; + } else { + [m appendFormat:@" fp16 %@_invd = const()[name=string(\"%@_invd\"), val=fp16(%f)];\n", p, p, inv_dim]; + [m appendFormat:@" tensor %@_ms = mul(x=%@_ss, y=%@_invd)[name=string(\"%@_ms\")];\n", + seq, p, p, p, p]; + + // rsqrt(ms + eps) + [m appendFormat:@" fp16 %@_eps = const()[name=string(\"%@_eps\"), val=fp16(%f)];\n", p, p, eps]; + [m appendFormat:@" tensor %@_mse = add(x=%@_ms, y=%@_eps)[name=string(\"%@_mse\")];\n", + seq, p, p, p, p]; + [m appendFormat:@" fp16 %@_nhalf = const()[name=string(\"%@_nhalf\"), val=fp16(-0.5)];\n", p, p]; + [m appendFormat:@" tensor %@_rrms = pow(x=%@_mse, y=%@_nhalf)[name=string(\"%@_rrms\")];\n", + seq, p, p, p, p]; + + // x * rrms + [m appendFormat:@" tensor %@_xr = mul(x=%@, y=%@_rrms)[name=string(\"%@_xr\")];\n", + dim, seq, p, inp, p, p]; + } // weight * normalized [m appendFormat:@" tensor %@_w = const()[name=string(\"%@_w\"), " @@ -260,9 +293,17 @@ NSString *inp = @(input); NSMutableString *m = [NSMutableString string]; - // SiLU(x) = x * sigmoid(x) - [m appendFormat:@" tensor %@_sig = sigmoid(x=%@)[name=string(\"%@_sig\")];\n", - dim, seq, p, inp, p]; + // SiLU(x) = x * sigmoid(x), using sigmoid(x) = 0.5 * (tanh(0.5 * x) + 1) + [m appendFormat:@" fp16 %@_half = const()[name=string(\"%@_half\"), val=fp16(0.5)];\n", p, p]; + [m appendFormat:@" tensor %@_hx = mul(x=%@, y=%@_half)[name=string(\"%@_hx\")];\n", + dim, seq, p, inp, p, p]; + [m appendFormat:@" tensor %@_th = tanh(x=%@_hx)[name=string(\"%@_th\")];\n", + dim, seq, p, p, p]; + [m appendFormat:@" fp16 %@_one = const()[name=string(\"%@_one\"), val=fp16(1.0)];\n", p, p]; + [m appendFormat:@" tensor %@_onep = add(x=%@_th, y=%@_one)[name=string(\"%@_onep\")];\n", + dim, seq, p, p, p, p]; + [m appendFormat:@" tensor %@_sig = mul(x=%@_onep, y=%@_half)[name=string(\"%@_sig\")];\n", + dim, seq, p, p, p, p]; [m appendFormat:@" tensor %@_out = mul(x=%@, y=%@_sig)[name=string(\"%@_out\")];\n", dim, seq, p, inp, p, p]; diff --git a/core/model_registry.m b/core/model_registry.m index 7553229..4c8551d 100644 --- a/core/model_registry.m +++ b/core/model_registry.m @@ -15,6 +15,7 @@ .hidden_dim = 3072, .vocab = 50257, .max_seq = 1024, + .n_kv_head = 12, }, .buckets = kGPT2RegistryBuckets, .n_buckets = 6, @@ -30,11 +31,44 @@ .hidden_dim = 2048, .vocab = 32000, .max_seq = 256, + .n_kv_head = 12, }, .buckets = NULL, .n_buckets = 0, .default_weights_dir = "model/blobs/stories110m", }, + { + .name = "qwen35_08b", + .config = { + .n_layer = 24, + .n_head = 8, + .d_model = 1024, + .head_dim = 256, + .hidden_dim = 3584, + .vocab = 248320, + .max_seq = 262144, + .n_kv_head = 2, + }, + .buckets = NULL, + .n_buckets = 0, + .default_weights_dir = "model/blobs/qwen35_08b", + }, + { + .name = "qwen35_9b", + .config = { + .n_layer = 32, + .n_head = 16, + .d_model = 4096, + .head_dim = 256, + .hidden_dim = 12288, + .vocab = 248320, + .max_seq = 262144, + .n_kv_head = 4, + }, + .buckets = NULL, + .n_buckets = 0, + .default_weights_dir = "model/blobs/qwen35_9b", + }, }; static const int kModelCount = sizeof(kModels) / sizeof(kModels[0]); diff --git a/docs/orion_q/ADR-008-orion-q-boundary.md b/docs/orion_q/ADR-008-orion-q-boundary.md new file mode 100644 index 0000000..2523c59 --- /dev/null +++ b/docs/orion_q/ADR-008-orion-q-boundary.md @@ -0,0 +1,45 @@ +# ADR-008: Orion-Q Boundary + +**Status**: Accepted +**Date**: 2026-03-15 + +## Decision + +This repository uses the following boundary: + +- `Orion`: the model-agnostic execution core +- `Orion-Q`: the Qwen-focused subset built on top of Orion + +`Orion-Q` is not a separate engine. It is the part of the Orion worktree that makes Qwen-family models runnable, testable, and trainable inside Orion. + +## Orion Includes + +- shared compiler and runtime +- shared model registry +- shared tokenizer and weight loading base +- CPU and ANE execution primitives + +## Orion-Q Includes + +- Qwen frontend and model registration +- Qwen weight loading and export path +- Qwen CPU inference path +- Qwen ANE or hybrid inference path +- Qwen LoRA training primitives +- Qwen diagnostics: smoke, probe, parity, and diff tests + +## Orion-Q Excludes + +- downstream training tracks +- domain-specific assets +- generated local reports and logs +- exported checkpoints and blobs +- Silver accelerator work + +## Result + +When sharing Orion-Q, it should be framed as: + +`a Qwen-focused subset built on top of Orion` + +and not as a disconnected new engine. diff --git a/docs/orion_q/ORION_Q_ANE_PREP_CLOSEOUT.md b/docs/orion_q/ORION_Q_ANE_PREP_CLOSEOUT.md new file mode 100644 index 0000000..e575f95 --- /dev/null +++ b/docs/orion_q/ORION_Q_ANE_PREP_CLOSEOUT.md @@ -0,0 +1,29 @@ +# Orion-Q ANE Training Preparation Close-Out + +## Summary + +The ANE training preparation line for Orion-Q is considered complete for mainline-candidate readiness. + +Validated path: + +- `ANE q/v base` +- `CPU LoRA delta / backward / update` + +## Validated Stages + +- forward probe +- smoke1 +- smoke10 +- micro-canary +- preflight + +## Meaning + +This close-out means Orion-Q has a validated ANE-assisted training candidate path. + +It does not mean: + +- full end-to-end ANE-only training +- performance superiority over CPU has already been proven + +Those remain separate optimization or promotion questions. diff --git a/docs/orion_q/ORION_Q_HYBRID_PARITY_CLOSEOUT.md b/docs/orion_q/ORION_Q_HYBRID_PARITY_CLOSEOUT.md new file mode 100644 index 0000000..7f28b84 --- /dev/null +++ b/docs/orion_q/ORION_Q_HYBRID_PARITY_CLOSEOUT.md @@ -0,0 +1,35 @@ +# Orion-Q Hybrid Parity Close-Out + +## Summary + +Hybrid parity for the current Orion-Q target smoke scope is closed. + +Target prompt set: + +- `사진` +- `NO` +- `안녕하세요` +- `정답은` + +Closed modes: + +- `single` +- `all_full` + +## Final State + +For the target prompt set above: + +- `single exact parity`: closed +- `all_full exact parity`: closed + +This close-out uses the current Orion-Q production path and diagnostics, not a disconnected external harness. + +## Validation Meaning + +The parity close-out means: + +- no remaining top-1 divergence in the target smoke scope +- diagnostics and parity checks can be treated as closed for Orion-Q scope + +It does not claim broader downstream task quality by itself. diff --git a/docs/orion_q/ORION_Q_PORT_CLOSEOUT.md b/docs/orion_q/ORION_Q_PORT_CLOSEOUT.md new file mode 100644 index 0000000..19e6bfd --- /dev/null +++ b/docs/orion_q/ORION_Q_PORT_CLOSEOUT.md @@ -0,0 +1,27 @@ +# Orion-Q Port Close-Out + +## Summary + +The Qwen porting scope for Orion-Q is considered closed for the currently defined target. + +Closed areas: + +- Qwen config support +- tokenizer integration +- weight loading and blob conversion path +- CPU inference path +- ANE prefill and hybrid execution path +- LoRA-oriented training primitives + +## Interpretation + +Within the Orion-Q boundary, the port is no longer an exploratory spike. +It is a working Qwen execution and validation subset on top of Orion. + +## Remaining Work Outside This Close-Out + +The following are intentionally outside the port close-out itself: + +- downstream domain tracks +- Silver accelerator work +- performance tuning for separate training programs diff --git a/kernels/inference/qwen_cpu_ops.h b/kernels/inference/qwen_cpu_ops.h new file mode 100644 index 0000000..4d0c8ed --- /dev/null +++ b/kernels/inference/qwen_cpu_ops.h @@ -0,0 +1,158 @@ +#ifndef ORION_QWEN_CPU_OPS_H +#define ORION_QWEN_CPU_OPS_H + +/// Qwen RMSNorm without bias. +/// out = x / sqrt(mean(x^2) + eps) * weight +void orion_qwen_cpu_rmsnorm(const float* x, const float* weight, int dim, float eps, float* out); + +/// Qwen SwiGLU MLP block: +/// hidden = silu(x @ gate_proj^T) * (x @ up_proj^T) +/// out = hidden @ down_proj^T +/// +/// Weights are expected in Orion blob layout: +/// - gate_proj: [d_ff, d_model] +/// - up_proj: [d_ff, d_model] +/// - down_proj: [d_model, d_ff] +void orion_qwen_cpu_swiglu_ffn(const float* x, + const float* gate_proj, + const float* up_proj, + const float* down_proj, + int d_model, + int d_ff, + float* out); + +/// Qwen3.5 full-attention CPU reference path without RoPE. +/// This follows the actual Qwen3.5 projection semantics: +/// - q_proj: [2 * (n_head * head_dim), d_model] +/// - query, gate = split(q_proj(x), 2) +/// - k_proj / v_proj: [n_kv_head * head_dim, d_model] +/// - o_proj: [d_model, n_head * head_dim] +/// - query/key use per-head RMSNorm; post-attention applies sigmoid(gate) +/// This is still a staging implementation and intentionally skips rotary. +/// +/// Inputs: +/// - x_seq: [seq_len, d_model] +/// - q_proj: [2 * (n_head * head_dim), d_model] +/// - k_proj: [n_kv_head * head_dim, d_model] +/// - v_proj: [n_kv_head * head_dim, d_model] +/// - o_proj: [d_model, n_head * head_dim] +/// - q_norm / k_norm: [head_dim] +/// +/// Output: +/// - out_seq: [seq_len, d_model] +void orion_qwen_cpu_full_attention_prefill_no_rope(const float* x_seq, + int seq_len, + const float* q_proj, + const float* k_proj, + const float* v_proj, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + float* out_seq); + +/// Qwen3.5 full-attention CPU reference path with text-only RoPE. +/// For text-only usage, the three MRoPE axes share the same scalar token position, +/// so standard RoPE over the partial rotary sub-dimension is sufficient. +void orion_qwen_cpu_full_attention_prefill_with_rope(const float* x_seq, + int seq_len, + const float* q_proj, + const float* k_proj, + const float* v_proj, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + float* out_seq); + +/// Qwen3.5 full-attention CPU reference path from precomputed projection outputs. +/// Inputs are the direct outputs of the ANE prefill q/kv graphs, transposed back +/// to CPU row-major layout: +/// - q_proj_out_seq: [seq_len, 2 * (n_head * head_dim)] +/// - k_proj_out_seq: [seq_len, n_kv_head * head_dim] +/// - v_proj_out_seq: [seq_len, n_kv_head * head_dim] +/// +/// This keeps the post-projection semantics identical to the regular CPU path: +/// per-head RMSNorm, partial RoPE, grouped-query attention, sigmoid gate, o_proj. +void orion_qwen_cpu_full_attention_from_projections_with_rope( + const float* q_proj_out_seq, + const float* k_proj_out_seq, + const float* v_proj_out_seq, + int seq_len, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + float* out_seq); + +/// Qwen3.5 linear-attention CPU preparation path. +/// This implements the front half of Qwen3.5GatedDeltaNet: +/// - in_proj_qkv / in_proj_z / in_proj_a / in_proj_b +/// - depthwise causal conv1d + silu +/// - split into q/k/v +/// - beta = sigmoid(b) +/// - g = -exp(A_log) * softplus(a + dt_bias) +/// +/// It intentionally stops before the gated delta recurrent core. +void orion_qwen_cpu_linear_attention_prep(const float* x_seq, + int seq_len, + const float* in_proj_qkv, + const float* in_proj_z, + const float* in_proj_a, + const float* in_proj_b, + const float* conv1d, + const float* dt_bias, + const float* a_log, + int d_model, + int num_k_heads, + int num_v_heads, + int head_k_dim, + int head_v_dim, + int conv_kernel, + float* query_out, + float* key_out, + float* value_out, + float* z_out, + float* beta_out, + float* g_out); + +/// Qwen3.5 linear-attention CPU reference path. +/// This extends the prep path with: +/// - l2-normalized recurrent gated delta core +/// - per-head RMSNormGated using z as gate input +/// - out_proj back to d_model +/// +/// It intentionally omits cache handling and only targets short prefill smoke. +void orion_qwen_cpu_linear_attention_recurrent_prefill(const float* x_seq, + int seq_len, + const float* in_proj_qkv, + const float* in_proj_z, + const float* in_proj_a, + const float* in_proj_b, + const float* conv1d, + const float* dt_bias, + const float* a_log, + const float* norm_weight, + const float* out_proj, + int d_model, + int num_k_heads, + int num_v_heads, + int head_k_dim, + int head_v_dim, + int conv_kernel, + float* out_seq); + +#endif // ORION_QWEN_CPU_OPS_H diff --git a/kernels/inference/qwen_cpu_ops.m b/kernels/inference/qwen_cpu_ops.m new file mode 100644 index 0000000..d17ad44 --- /dev/null +++ b/kernels/inference/qwen_cpu_ops.m @@ -0,0 +1,758 @@ +#import "qwen_cpu_ops.h" +#import +#import +#import + +void orion_qwen_cpu_rmsnorm(const float* x, const float* weight, int dim, float eps, float* out) { + float mean_sq = 0.0f; + vDSP_measqv(x, 1, &mean_sq, dim); + float scale = 1.0f / sqrtf(mean_sq + eps); + vDSP_vsmul(x, 1, &scale, out, 1, dim); + vDSP_vmul(out, 1, weight, 1, out, 1, dim); +} + +static void linear_no_bias(const float* x, const float* w, int in_dim, int out_dim, float* out) { + cblas_sgemv(CblasRowMajor, CblasNoTrans, + out_dim, in_dim, + 1.0f, w, in_dim, x, 1, + 0.0f, out, 1); +} + +static inline float silu(float x) { + return x / (1.0f + expf(-x)); +} + +static inline float sigmoidf_safe(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +static inline float softplusf_safe(float x) { + if (x > 20.0f) return x; + return log1pf(expf(x)); +} + +static void l2norm_vec(const float* x, int dim, float eps, float* out) { + float sumsq = 0.0f; + for (int i = 0; i < dim; i++) { + sumsq += x[i] * x[i]; + } + float inv = 1.0f / sqrtf(sumsq + eps); + for (int i = 0; i < dim; i++) { + out[i] = x[i] * inv; + } +} + +static void qwen_rmsnorm_gated(const float* x, const float* gate, const float* weight, int dim, float eps, float* out) { + float mean_sq = 0.0f; + for (int i = 0; i < dim; i++) { + mean_sq += x[i] * x[i]; + } + mean_sq /= (float)dim; + float inv = 1.0f / sqrtf(mean_sq + eps); + for (int i = 0; i < dim; i++) { + out[i] = (x[i] * inv) * weight[i] * silu(gate[i]); + } +} + +static void apply_rope_text_inplace(float* q, + float* k, + int seq_len, + int n_q_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor) { + int rotary_dim = (int)(head_dim * partial_rotary_factor); + if (rotary_dim > head_dim) rotary_dim = head_dim; + if (rotary_dim % 2 != 0) rotary_dim -= 1; + if (rotary_dim <= 0) return; + + int half_rot = rotary_dim / 2; + float* inv_freq = (float*)malloc((size_t)half_rot * sizeof(float)); + if (!inv_freq) return; + + for (int i = 0; i < half_rot; i++) { + float exponent = (2.0f * (float)i) / (float)rotary_dim; + inv_freq[i] = 1.0f / powf(rope_theta, exponent); + } + + for (int pos = 0; pos < seq_len; pos++) { + for (int h = 0; h < n_q_head; h++) { + float* qh = q + pos * (n_q_head * head_dim) + h * head_dim; + for (int i = 0; i < half_rot; i++) { + float angle = (float)pos * inv_freq[i]; + float c = cosf(angle); + float s = sinf(angle); + float x0 = qh[i]; + float x1 = qh[i + half_rot]; + qh[i] = x0 * c - x1 * s; + qh[i + half_rot] = x1 * c + x0 * s; + } + } + for (int h = 0; h < n_kv_head; h++) { + float* kh = k + pos * (n_kv_head * head_dim) + h * head_dim; + for (int i = 0; i < half_rot; i++) { + float angle = (float)pos * inv_freq[i]; + float c = cosf(angle); + float s = sinf(angle); + float x0 = kh[i]; + float x1 = kh[i + half_rot]; + kh[i] = x0 * c - x1 * s; + kh[i + half_rot] = x1 * c + x0 * s; + } + } + } + + free(inv_freq); +} + +static void qwen_full_attention_from_projected_qkv(const float* q_proj_out_seq, + const float* k_proj_out_seq, + const float* v_proj_out_seq, + int seq_len, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + int apply_rope, + float rope_theta, + float partial_rotary_factor, + float* out_seq) { + int q_dim = n_head * head_dim; + int kv_dim = n_kv_head * head_dim; + int q_per_kv = n_head / n_kv_head; + float scale = 1.0f / sqrtf((float)head_dim); + + float* q = (float*)malloc((size_t)seq_len * q_dim * sizeof(float)); + float* gate = (float*)malloc((size_t)seq_len * q_dim * sizeof(float)); + float* k = (float*)malloc((size_t)seq_len * kv_dim * sizeof(float)); + float* v = (float*)malloc((size_t)seq_len * kv_dim * sizeof(float)); + float* attn_cat = (float*)calloc((size_t)seq_len * q_dim, sizeof(float)); + float* scores = (float*)malloc((size_t)seq_len * seq_len * sizeof(float)); + + memcpy(k, k_proj_out_seq, (size_t)seq_len * kv_dim * sizeof(float)); + memcpy(v, v_proj_out_seq, (size_t)seq_len * kv_dim * sizeof(float)); + for (int s = 0; s < seq_len; s++) { + memcpy(q + s * q_dim, + q_proj_out_seq + s * (q_dim * 2), + (size_t)q_dim * sizeof(float)); + memcpy(gate + s * q_dim, + q_proj_out_seq + s * (q_dim * 2) + q_dim, + (size_t)q_dim * sizeof(float)); + } + + for (int s = 0; s < seq_len; s++) { + for (int h = 0; h < n_head; h++) { + orion_qwen_cpu_rmsnorm(q + s * q_dim + h * head_dim, q_norm, head_dim, 1e-6f, + q + s * q_dim + h * head_dim); + } + for (int h = 0; h < n_kv_head; h++) { + orion_qwen_cpu_rmsnorm(k + s * kv_dim + h * head_dim, k_norm, head_dim, 1e-6f, + k + s * kv_dim + h * head_dim); + } + } + + if (apply_rope) { + apply_rope_text_inplace(q, k, seq_len, n_head, n_kv_head, head_dim, rope_theta, partial_rotary_factor); + } + + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + float* qh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* kh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* vh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* attn_h = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + memcpy(qh + s * head_dim, q + s * q_dim + h * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(kh + s * head_dim, k + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(vh + s * head_dim, v + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, seq_len, head_dim, + scale, qh, head_dim, kh, head_dim, + 0.0f, scores, seq_len); + + for (int i = 0; i < seq_len; i++) { + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = -INFINITY; + } + float max_val = scores[i * seq_len]; + for (int j = 1; j <= i; j++) { + if (scores[i * seq_len + j] > max_val) { + max_val = scores[i * seq_len + j]; + } + } + float sum = 0.0f; + for (int j = 0; j <= i; j++) { + scores[i * seq_len + j] = expf(scores[i * seq_len + j] - max_val); + sum += scores[i * seq_len + j]; + } + for (int j = 0; j <= i; j++) { + scores[i * seq_len + j] /= sum; + } + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = 0.0f; + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + seq_len, head_dim, seq_len, + 1.0f, scores, seq_len, vh, head_dim, + 0.0f, attn_h, head_dim); + + for (int s = 0; s < seq_len; s++) { + memcpy(attn_cat + s * q_dim + h * head_dim, + attn_h + s * head_dim, + (size_t)head_dim * sizeof(float)); + } + + free(qh); + free(kh); + free(vh); + free(attn_h); + } + + for (int i = 0; i < seq_len * q_dim; i++) { + attn_cat[i] *= sigmoidf_safe(gate[i]); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, d_model, q_dim, + 1.0f, attn_cat, q_dim, o_proj, q_dim, + 0.0f, out_seq, d_model); + + free(q); + free(gate); + free(k); + free(v); + free(attn_cat); + free(scores); +} + +void orion_qwen_cpu_swiglu_ffn(const float* x, + const float* gate_proj, + const float* up_proj, + const float* down_proj, + int d_model, + int d_ff, + float* out) { + float* gate = (float*)malloc((size_t)d_ff * sizeof(float)); + float* up = (float*)malloc((size_t)d_ff * sizeof(float)); + float* hidden = (float*)malloc((size_t)d_ff * sizeof(float)); + + linear_no_bias(x, gate_proj, d_model, d_ff, gate); + linear_no_bias(x, up_proj, d_model, d_ff, up); + + for (int i = 0; i < d_ff; i++) { + hidden[i] = silu(gate[i]) * up[i]; + } + + linear_no_bias(hidden, down_proj, d_ff, d_model, out); + + free(gate); + free(up); + free(hidden); +} + +void orion_qwen_cpu_full_attention_prefill_no_rope(const float* x_seq, + int seq_len, + const float* q_proj, + const float* k_proj, + const float* v_proj, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + float* out_seq) { + int q_dim = n_head * head_dim; + int kv_dim = n_kv_head * head_dim; + int q_per_kv = n_head / n_kv_head; + float scale = 1.0f / sqrtf((float)head_dim); + + float* q_full = (float*)malloc((size_t)seq_len * (q_dim * 2) * sizeof(float)); + float* q = (float*)malloc((size_t)seq_len * q_dim * sizeof(float)); + float* gate = (float*)malloc((size_t)seq_len * q_dim * sizeof(float)); + float* k = (float*)malloc((size_t)seq_len * kv_dim * sizeof(float)); + float* v = (float*)malloc((size_t)seq_len * kv_dim * sizeof(float)); + float* attn_cat = (float*)calloc((size_t)seq_len * q_dim, sizeof(float)); + float* scores = (float*)malloc((size_t)seq_len * seq_len * sizeof(float)); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, q_dim * 2, d_model, + 1.0f, x_seq, d_model, q_proj, d_model, + 0.0f, q_full, q_dim * 2); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, kv_dim, d_model, + 1.0f, x_seq, d_model, k_proj, d_model, + 0.0f, k, kv_dim); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, kv_dim, d_model, + 1.0f, x_seq, d_model, v_proj, d_model, + 0.0f, v, kv_dim); + + for (int s = 0; s < seq_len; s++) { + memcpy(q + s * q_dim, q_full + s * (q_dim * 2), (size_t)q_dim * sizeof(float)); + memcpy(gate + s * q_dim, q_full + s * (q_dim * 2) + q_dim, (size_t)q_dim * sizeof(float)); + } + + for (int s = 0; s < seq_len; s++) { + for (int h = 0; h < n_head; h++) { + orion_qwen_cpu_rmsnorm(q + s * q_dim + h * head_dim, q_norm, head_dim, 1e-6f, + q + s * q_dim + h * head_dim); + } + for (int h = 0; h < n_kv_head; h++) { + orion_qwen_cpu_rmsnorm(k + s * kv_dim + h * head_dim, k_norm, head_dim, 1e-6f, + k + s * kv_dim + h * head_dim); + } + } + + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + float* qh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* kh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* vh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* attn_h = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + memcpy(qh + s * head_dim, q + s * q_dim + h * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(kh + s * head_dim, k + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(vh + s * head_dim, v + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, seq_len, head_dim, + scale, qh, head_dim, kh, head_dim, + 0.0f, scores, seq_len); + + for (int i = 0; i < seq_len; i++) { + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = -INFINITY; + } + float max_val = scores[i * seq_len]; + for (int j = 1; j <= i; j++) { + if (scores[i * seq_len + j] > max_val) { + max_val = scores[i * seq_len + j]; + } + } + float sum = 0.0f; + for (int j = 0; j <= i; j++) { + scores[i * seq_len + j] = expf(scores[i * seq_len + j] - max_val); + sum += scores[i * seq_len + j]; + } + for (int j = 0; j <= i; j++) { + scores[i * seq_len + j] /= sum; + } + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = 0.0f; + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + seq_len, head_dim, seq_len, + 1.0f, scores, seq_len, vh, head_dim, + 0.0f, attn_h, head_dim); + + for (int s = 0; s < seq_len; s++) { + memcpy(attn_cat + s * q_dim + h * head_dim, attn_h + s * head_dim, + (size_t)head_dim * sizeof(float)); + } + + free(qh); + free(kh); + free(vh); + free(attn_h); + } + + for (int i = 0; i < seq_len * q_dim; i++) { + attn_cat[i] *= sigmoidf_safe(gate[i]); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, d_model, q_dim, + 1.0f, attn_cat, q_dim, o_proj, q_dim, + 0.0f, out_seq, d_model); + + free(q_full); + free(q); + free(gate); + free(k); + free(v); + free(attn_cat); + free(scores); +} + +void orion_qwen_cpu_full_attention_prefill_with_rope(const float* x_seq, + int seq_len, + const float* q_proj, + const float* k_proj, + const float* v_proj, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + float* out_seq) { + int q_dim = n_head * head_dim; + int kv_dim = n_kv_head * head_dim; + int q_per_kv = n_head / n_kv_head; + float scale = 1.0f / sqrtf((float)head_dim); + + float* q_full = (float*)malloc((size_t)seq_len * (q_dim * 2) * sizeof(float)); + float* q = (float*)malloc((size_t)seq_len * q_dim * sizeof(float)); + float* gate = (float*)malloc((size_t)seq_len * q_dim * sizeof(float)); + float* k = (float*)malloc((size_t)seq_len * kv_dim * sizeof(float)); + float* v = (float*)malloc((size_t)seq_len * kv_dim * sizeof(float)); + float* attn_cat = (float*)calloc((size_t)seq_len * q_dim, sizeof(float)); + float* scores = (float*)malloc((size_t)seq_len * seq_len * sizeof(float)); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, q_dim * 2, d_model, + 1.0f, x_seq, d_model, q_proj, d_model, + 0.0f, q_full, q_dim * 2); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, kv_dim, d_model, + 1.0f, x_seq, d_model, k_proj, d_model, + 0.0f, k, kv_dim); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, kv_dim, d_model, + 1.0f, x_seq, d_model, v_proj, d_model, + 0.0f, v, kv_dim); + + for (int s = 0; s < seq_len; s++) { + memcpy(q + s * q_dim, q_full + s * (q_dim * 2), (size_t)q_dim * sizeof(float)); + memcpy(gate + s * q_dim, q_full + s * (q_dim * 2) + q_dim, (size_t)q_dim * sizeof(float)); + } + + for (int s = 0; s < seq_len; s++) { + for (int h = 0; h < n_head; h++) { + orion_qwen_cpu_rmsnorm(q + s * q_dim + h * head_dim, q_norm, head_dim, 1e-6f, + q + s * q_dim + h * head_dim); + } + for (int h = 0; h < n_kv_head; h++) { + orion_qwen_cpu_rmsnorm(k + s * kv_dim + h * head_dim, k_norm, head_dim, 1e-6f, + k + s * kv_dim + h * head_dim); + } + } + + apply_rope_text_inplace(q, k, seq_len, n_head, n_kv_head, head_dim, rope_theta, partial_rotary_factor); + + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + float* qh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* kh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* vh = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + float* attn_h = (float*)malloc((size_t)seq_len * head_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + memcpy(qh + s * head_dim, q + s * q_dim + h * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(kh + s * head_dim, k + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(vh + s * head_dim, v + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, seq_len, head_dim, + scale, qh, head_dim, kh, head_dim, + 0.0f, scores, seq_len); + + for (int i = 0; i < seq_len; i++) { + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = -INFINITY; + } + float max_val = scores[i * seq_len]; + for (int j = 1; j <= i; j++) { + if (scores[i * seq_len + j] > max_val) { + max_val = scores[i * seq_len + j]; + } + } + float sum = 0.0f; + for (int j = 0; j <= i; j++) { + scores[i * seq_len + j] = expf(scores[i * seq_len + j] - max_val); + sum += scores[i * seq_len + j]; + } + for (int j = 0; j <= i; j++) { + scores[i * seq_len + j] /= sum; + } + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = 0.0f; + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + seq_len, head_dim, seq_len, + 1.0f, scores, seq_len, vh, head_dim, + 0.0f, attn_h, head_dim); + + for (int s = 0; s < seq_len; s++) { + memcpy(attn_cat + s * q_dim + h * head_dim, attn_h + s * head_dim, + (size_t)head_dim * sizeof(float)); + } + + free(qh); + free(kh); + free(vh); + free(attn_h); + } + + for (int i = 0; i < seq_len * q_dim; i++) { + attn_cat[i] *= sigmoidf_safe(gate[i]); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, d_model, q_dim, + 1.0f, attn_cat, q_dim, o_proj, q_dim, + 0.0f, out_seq, d_model); + + free(q_full); + free(q); + free(gate); + free(k); + free(v); + free(attn_cat); + free(scores); +} + +void orion_qwen_cpu_linear_attention_prep(const float* x_seq, + int seq_len, + const float* in_proj_qkv, + const float* in_proj_z, + const float* in_proj_a, + const float* in_proj_b, + const float* conv1d, + const float* dt_bias, + const float* a_log, + int d_model, + int num_k_heads, + int num_v_heads, + int head_k_dim, + int head_v_dim, + int conv_kernel, + float* query_out, + float* key_out, + float* value_out, + float* z_out, + float* beta_out, + float* g_out) { + int key_dim = num_k_heads * head_k_dim; + int value_dim = num_v_heads * head_v_dim; + int conv_dim = key_dim * 2 + value_dim; + + float* mixed = (float*)malloc((size_t)seq_len * conv_dim * sizeof(float)); + float* mixed_conv = (float*)malloc((size_t)seq_len * conv_dim * sizeof(float)); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, conv_dim, d_model, + 1.0f, x_seq, d_model, in_proj_qkv, d_model, + 0.0f, mixed, conv_dim); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, value_dim, d_model, + 1.0f, x_seq, d_model, in_proj_z, d_model, + 0.0f, z_out, value_dim); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, num_v_heads, d_model, + 1.0f, x_seq, d_model, in_proj_b, d_model, + 0.0f, beta_out, num_v_heads); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, num_v_heads, d_model, + 1.0f, x_seq, d_model, in_proj_a, d_model, + 0.0f, g_out, num_v_heads); + + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < conv_dim; c++) { + float sum = 0.0f; + const float* kernel = conv1d + c * conv_kernel; + for (int k = 0; k < conv_kernel; k++) { + int src_t = t - (conv_kernel - 1) + k; + float x = 0.0f; + if (src_t >= 0 && src_t < seq_len) { + x = mixed[src_t * conv_dim + c]; + } + sum += kernel[k] * x; + } + mixed_conv[t * conv_dim + c] = silu(sum); + } + } + + for (int t = 0; t < seq_len; t++) { + const float* row = mixed_conv + t * conv_dim; + memcpy(query_out + t * key_dim, row, (size_t)key_dim * sizeof(float)); + memcpy(key_out + t * key_dim, row + key_dim, (size_t)key_dim * sizeof(float)); + memcpy(value_out + t * value_dim, row + key_dim * 2, (size_t)value_dim * sizeof(float)); + } + + for (int i = 0; i < seq_len * num_v_heads; i++) { + beta_out[i] = sigmoidf_safe(beta_out[i]); + } + + for (int t = 0; t < seq_len; t++) { + for (int h = 0; h < num_v_heads; h++) { + float a = g_out[t * num_v_heads + h]; + float dt = dt_bias[h]; + float al = a_log[h]; + g_out[t * num_v_heads + h] = -expf(al) * softplusf_safe(a + dt); + } + } + + free(mixed); + free(mixed_conv); +} + +void orion_qwen_cpu_linear_attention_recurrent_prefill(const float* x_seq, + int seq_len, + const float* in_proj_qkv, + const float* in_proj_z, + const float* in_proj_a, + const float* in_proj_b, + const float* conv1d, + const float* dt_bias, + const float* a_log, + const float* norm_weight, + const float* out_proj, + int d_model, + int num_k_heads, + int num_v_heads, + int head_k_dim, + int head_v_dim, + int conv_kernel, + float* out_seq) { + int key_dim = num_k_heads * head_k_dim; + int value_dim = num_v_heads * head_v_dim; + float scale = 1.0f / sqrtf((float)head_k_dim); + + float* query = (float*)calloc((size_t)seq_len * key_dim, sizeof(float)); + float* key = (float*)calloc((size_t)seq_len * key_dim, sizeof(float)); + float* value = (float*)calloc((size_t)seq_len * value_dim, sizeof(float)); + float* z = (float*)calloc((size_t)seq_len * value_dim, sizeof(float)); + float* beta = (float*)calloc((size_t)seq_len * num_v_heads, sizeof(float)); + float* g = (float*)calloc((size_t)seq_len * num_v_heads, sizeof(float)); + float* core = (float*)calloc((size_t)seq_len * value_dim, sizeof(float)); + float* gated = (float*)calloc((size_t)seq_len * value_dim, sizeof(float)); + float* state = (float*)calloc((size_t)num_v_heads * head_k_dim * head_v_dim, sizeof(float)); + float* q_norm = (float*)malloc((size_t)head_k_dim * sizeof(float)); + float* k_norm = (float*)malloc((size_t)head_k_dim * sizeof(float)); + float* kv_mem = (float*)malloc((size_t)head_v_dim * sizeof(float)); + float* delta = (float*)malloc((size_t)head_v_dim * sizeof(float)); + float* core_head = (float*)malloc((size_t)head_v_dim * sizeof(float)); + + orion_qwen_cpu_linear_attention_prep( + x_seq, seq_len, + in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, conv1d, dt_bias, a_log, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + query, key, value, z, beta, g + ); + + for (int t = 0; t < seq_len; t++) { + for (int h = 0; h < num_v_heads; h++) { + const float* q_t = query + t * key_dim + h * head_k_dim; + const float* k_t = key + t * key_dim + h * head_k_dim; + const float* v_t = value + t * value_dim + h * head_v_dim; + float* state_h = state + h * head_k_dim * head_v_dim; + float* core_t = core + t * value_dim + h * head_v_dim; + const float* z_t = z + t * value_dim + h * head_v_dim; + + l2norm_vec(q_t, head_k_dim, 1e-6f, q_norm); + l2norm_vec(k_t, head_k_dim, 1e-6f, k_norm); + for (int i = 0; i < head_k_dim; i++) { + q_norm[i] *= scale; + } + + float decay = expf(g[t * num_v_heads + h]); + float beta_t = beta[t * num_v_heads + h]; + + for (int i = 0; i < head_k_dim * head_v_dim; i++) { + state_h[i] *= decay; + } + + for (int v_idx = 0; v_idx < head_v_dim; v_idx++) { + float sum = 0.0f; + for (int k_idx = 0; k_idx < head_k_dim; k_idx++) { + sum += state_h[k_idx * head_v_dim + v_idx] * k_norm[k_idx]; + } + kv_mem[v_idx] = sum; + delta[v_idx] = (v_t[v_idx] - sum) * beta_t; + } + + for (int k_idx = 0; k_idx < head_k_dim; k_idx++) { + for (int v_idx = 0; v_idx < head_v_dim; v_idx++) { + state_h[k_idx * head_v_dim + v_idx] += k_norm[k_idx] * delta[v_idx]; + } + } + + for (int v_idx = 0; v_idx < head_v_dim; v_idx++) { + float sum = 0.0f; + for (int k_idx = 0; k_idx < head_k_dim; k_idx++) { + sum += state_h[k_idx * head_v_dim + v_idx] * q_norm[k_idx]; + } + core_head[v_idx] = sum; + } + + qwen_rmsnorm_gated(core_head, z_t, norm_weight, head_v_dim, 1e-6f, core_t); + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, d_model, value_dim, + 1.0f, core, value_dim, out_proj, value_dim, + 0.0f, out_seq, d_model); + + free(query); + free(key); + free(value); + free(z); + free(beta); + free(g); + free(core); + free(gated); + free(state); + free(q_norm); + free(k_norm); + free(kv_mem); + free(delta); + free(core_head); +} + +void orion_qwen_cpu_full_attention_from_projections_with_rope( + const float* q_proj_out_seq, + const float* k_proj_out_seq, + const float* v_proj_out_seq, + int seq_len, + const float* o_proj, + const float* q_norm, + const float* k_norm, + int d_model, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + float* out_seq) { + qwen_full_attention_from_projected_qkv(q_proj_out_seq, + k_proj_out_seq, + v_proj_out_seq, + seq_len, + o_proj, + q_norm, + k_norm, + d_model, + n_head, + n_kv_head, + head_dim, + 1, + rope_theta, + partial_rotary_factor, + out_seq); +} diff --git a/kernels/training/qwen_lora_cpu_ops.h b/kernels/training/qwen_lora_cpu_ops.h new file mode 100644 index 0000000..6778499 --- /dev/null +++ b/kernels/training/qwen_lora_cpu_ops.h @@ -0,0 +1,85 @@ +#ifndef ORION_QWEN_LORA_CPU_OPS_H +#define ORION_QWEN_LORA_CPU_OPS_H + +#import + +typedef struct OrionQwenStreamingCEContext OrionQwenStreamingCEContext; + +typedef struct { + int in_dim; + int out_dim; + int rank; + float alpha; + float scale; + float *a; + float *b; + float *da; + float *db; + float *ma; + float *va; + float *mb; + float *vb; +} OrionLoRAMatrix; + +void orion_qwen_lora_matrix_init(OrionLoRAMatrix *mat, + int in_dim, + int out_dim, + int rank, + float alpha, + unsigned int seed); + +void orion_qwen_lora_matrix_zero_grad(OrionLoRAMatrix *mat); + +void orion_qwen_lora_matrix_scale_grad(OrionLoRAMatrix *mat, float scale); + +void orion_qwen_lora_matrix_free(OrionLoRAMatrix *mat); + +void orion_qwen_lora_linear_forward(const float *x, + const float *w_base, + const OrionLoRAMatrix *lora, + float *y_out); + +void orion_qwen_lora_linear_delta_forward(const float *x, + const OrionLoRAMatrix *lora, + float *y_delta_out); + +void orion_qwen_lora_linear_backward(const float *x, + const float *w_base, + OrionLoRAMatrix *lora, + const float *dy, + float *dx_out); + +void orion_qwen_cpu_swiglu_ffn_bwd(const float *x, + const float *gate_proj, + const float *up_proj, + const float *down_proj, + int d_model, + int d_ff, + const float *dy_out, + float *dx_out); + +OrionQwenStreamingCEContext *orion_qwen_streaming_ce_context_open(const char *embed_blob_path, + int d_model, + int vocab); + +float orion_qwen_streaming_ce_tied_embedding_ctx(OrionQwenStreamingCEContext *ctx, + const float *hidden, + int target_token, + float *d_hidden_out); + +void orion_qwen_streaming_ce_context_close(OrionQwenStreamingCEContext *ctx); + +/// Streaming cross-entropy against tied embeddings. +/// Uses a 1-pass numerically stable accumulation over vocab rows. +float orion_qwen_cpu_streaming_ce_tied_embedding(const char *embed_blob_path, + const float *hidden, + int d_model, + int vocab, + int target_token, + float *d_hidden_out); + +double orion_qwen_lora_abs_sum(const float *x, int n); + +double orion_qwen_lora_grad_abs_sum(const OrionLoRAMatrix *mat); + +#endif // ORION_QWEN_LORA_CPU_OPS_H diff --git a/kernels/training/qwen_lora_cpu_ops.m b/kernels/training/qwen_lora_cpu_ops.m new file mode 100644 index 0000000..97e8ca3 --- /dev/null +++ b/kernels/training/qwen_lora_cpu_ops.m @@ -0,0 +1,312 @@ +#import "qwen_lora_cpu_ops.h" +#import "../inference/qwen_cpu_ops.h" +#import "../training/stories_cpu_ops.h" +#import "../../model/weight_loader.h" +#import +#import +#import +#import + +struct OrionQwenStreamingCEContext { + OrionBlobRowReader *reader; + int d_model; + int vocab; + float *row; + float *target; + float *expected; +}; + +static inline float lora_rand_unit(unsigned int *state) { + *state = (*state * 1664525u) + 1013904223u; + return ((float)((*state >> 8) & 0x00FFFFFF) / 16777215.0f) - 0.5f; +} + +static inline float silu_fwd(float x) { + return x / (1.0f + expf(-x)); +} + +static inline float silu_bwd(float x) { + float sig = 1.0f / (1.0f + expf(-x)); + return sig + x * sig * (1.0f - sig); +} + +double orion_qwen_lora_abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +double orion_qwen_lora_grad_abs_sum(const OrionLoRAMatrix *mat) { + if (!mat) return 0.0; + return orion_qwen_lora_abs_sum(mat->da, mat->rank * mat->in_dim) + + orion_qwen_lora_abs_sum(mat->db, mat->out_dim * mat->rank); +} + +void orion_qwen_lora_matrix_init(OrionLoRAMatrix *mat, + int in_dim, + int out_dim, + int rank, + float alpha, + unsigned int seed) { + memset(mat, 0, sizeof(*mat)); + mat->in_dim = in_dim; + mat->out_dim = out_dim; + mat->rank = rank; + mat->alpha = alpha; + mat->scale = alpha / (float)rank; + + size_t a_count = (size_t)rank * in_dim; + size_t b_count = (size_t)out_dim * rank; + mat->a = (float *)calloc(a_count, sizeof(float)); + mat->b = (float *)calloc(b_count, sizeof(float)); + mat->da = (float *)calloc(a_count, sizeof(float)); + mat->db = (float *)calloc(b_count, sizeof(float)); + mat->ma = (float *)calloc(a_count, sizeof(float)); + mat->va = (float *)calloc(a_count, sizeof(float)); + mat->mb = (float *)calloc(b_count, sizeof(float)); + mat->vb = (float *)calloc(b_count, sizeof(float)); + + unsigned int state = seed ? seed : 1u; + for (size_t i = 0; i < a_count; i++) { + mat->a[i] = 0.01f * lora_rand_unit(&state); + } +} + +void orion_qwen_lora_matrix_zero_grad(OrionLoRAMatrix *mat) { + if (!mat || !mat->da || !mat->db) return; + memset(mat->da, 0, (size_t)mat->rank * mat->in_dim * sizeof(float)); + memset(mat->db, 0, (size_t)mat->out_dim * mat->rank * sizeof(float)); +} + +void orion_qwen_lora_matrix_scale_grad(OrionLoRAMatrix *mat, float scale) { + if (!mat || !mat->da || !mat->db) return; + size_t a_count = (size_t)mat->rank * mat->in_dim; + size_t b_count = (size_t)mat->out_dim * mat->rank; + for (size_t i = 0; i < a_count; i++) mat->da[i] *= scale; + for (size_t i = 0; i < b_count; i++) mat->db[i] *= scale; +} + +void orion_qwen_lora_matrix_free(OrionLoRAMatrix *mat) { + if (!mat) return; + free(mat->a); free(mat->b); + free(mat->da); free(mat->db); + free(mat->ma); free(mat->va); + free(mat->mb); free(mat->vb); + memset(mat, 0, sizeof(*mat)); +} + +void orion_qwen_lora_linear_forward(const float *x, + const float *w_base, + const OrionLoRAMatrix *lora, + float *y_out) { + cblas_sgemv(CblasRowMajor, CblasNoTrans, + lora->out_dim, lora->in_dim, + 1.0f, w_base, lora->in_dim, x, 1, + 0.0f, y_out, 1); + + float *z = (float *)calloc((size_t)lora->rank, sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + lora->rank, lora->in_dim, + 1.0f, lora->a, lora->in_dim, x, 1, + 0.0f, z, 1); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + lora->out_dim, lora->rank, + lora->scale, lora->b, lora->rank, z, 1, + 1.0f, y_out, 1); + free(z); +} + +void orion_qwen_lora_linear_delta_forward(const float *x, + const OrionLoRAMatrix *lora, + float *y_delta_out) { + memset(y_delta_out, 0, (size_t)lora->out_dim * sizeof(float)); + float *z = (float *)calloc((size_t)lora->rank, sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + lora->rank, lora->in_dim, + 1.0f, lora->a, lora->in_dim, x, 1, + 0.0f, z, 1); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + lora->out_dim, lora->rank, + lora->scale, lora->b, lora->rank, z, 1, + 0.0f, y_delta_out, 1); + free(z); +} + +void orion_qwen_lora_linear_backward(const float *x, + const float *w_base, + OrionLoRAMatrix *lora, + const float *dy, + float *dx_out) { + memset(dx_out, 0, (size_t)lora->in_dim * sizeof(float)); + + cblas_sgemv(CblasRowMajor, CblasTrans, + lora->out_dim, lora->in_dim, + 1.0f, w_base, lora->in_dim, dy, 1, + 0.0f, dx_out, 1); + + float *z = (float *)calloc((size_t)lora->rank, sizeof(float)); + float *dz = (float *)calloc((size_t)lora->rank, sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + lora->rank, lora->in_dim, + 1.0f, lora->a, lora->in_dim, x, 1, + 0.0f, z, 1); + + for (int o = 0; o < lora->out_dim; o++) { + const float dy_o = dy[o]; + const float *b_row = lora->b + (size_t)o * lora->rank; + float *db_row = lora->db + (size_t)o * lora->rank; + for (int r = 0; r < lora->rank; r++) { + db_row[r] += lora->scale * dy_o * z[r]; + dz[r] += lora->scale * b_row[r] * dy_o; + } + } + + for (int r = 0; r < lora->rank; r++) { + float *da_row = lora->da + (size_t)r * lora->in_dim; + const float a_scale = dz[r]; + for (int i = 0; i < lora->in_dim; i++) { + da_row[i] += a_scale * x[i]; + dx_out[i] += lora->a[(size_t)r * lora->in_dim + i] * a_scale; + } + } + + free(z); + free(dz); +} + +void orion_qwen_cpu_swiglu_ffn_bwd(const float *x, + const float *gate_proj, + const float *up_proj, + const float *down_proj, + int d_model, + int d_ff, + const float *dy_out, + float *dx_out) { + float *gate = (float *)calloc((size_t)d_ff, sizeof(float)); + float *up = (float *)calloc((size_t)d_ff, sizeof(float)); + float *hidden = (float *)calloc((size_t)d_ff, sizeof(float)); + float *dhidden = (float *)calloc((size_t)d_ff, sizeof(float)); + float *dgate = (float *)calloc((size_t)d_ff, sizeof(float)); + float *dup = (float *)calloc((size_t)d_ff, sizeof(float)); + + cblas_sgemv(CblasRowMajor, CblasNoTrans, + d_ff, d_model, 1.0f, gate_proj, d_model, x, 1, 0.0f, gate, 1); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + d_ff, d_model, 1.0f, up_proj, d_model, x, 1, 0.0f, up, 1); + for (int i = 0; i < d_ff; i++) hidden[i] = silu_fwd(gate[i]) * up[i]; + + cblas_sgemv(CblasRowMajor, CblasTrans, + d_model, d_ff, 1.0f, down_proj, d_ff, dy_out, 1, 0.0f, dhidden, 1); + + for (int i = 0; i < d_ff; i++) { + dgate[i] = dhidden[i] * up[i] * silu_bwd(gate[i]); + dup[i] = dhidden[i] * silu_fwd(gate[i]); + } + + memset(dx_out, 0, (size_t)d_model * sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasTrans, + d_ff, d_model, 1.0f, gate_proj, d_model, dgate, 1, 1.0f, dx_out, 1); + cblas_sgemv(CblasRowMajor, CblasTrans, + d_ff, d_model, 1.0f, up_proj, d_model, dup, 1, 1.0f, dx_out, 1); + + free(gate); + free(up); + free(hidden); + free(dhidden); + free(dgate); + free(dup); +} + +float orion_qwen_cpu_streaming_ce_tied_embedding(const char *embed_blob_path, + const float *hidden, + int d_model, + int vocab, + int target_token, + float *d_hidden_out) { + OrionQwenStreamingCEContext *ctx = orion_qwen_streaming_ce_context_open(embed_blob_path, d_model, vocab); + if (!ctx) { + return NAN; + } + float loss = orion_qwen_streaming_ce_tied_embedding_ctx(ctx, hidden, target_token, d_hidden_out); + orion_qwen_streaming_ce_context_close(ctx); + return loss; +} + +OrionQwenStreamingCEContext *orion_qwen_streaming_ce_context_open(const char *embed_blob_path, + int d_model, + int vocab) { + if (!embed_blob_path || d_model <= 0 || vocab <= 0) return NULL; + OrionQwenStreamingCEContext *ctx = (OrionQwenStreamingCEContext *)calloc(1, sizeof(OrionQwenStreamingCEContext)); + if (!ctx) return NULL; + ctx->reader = orion_blob_row_reader_open(embed_blob_path, d_model); + ctx->d_model = d_model; + ctx->vocab = vocab; + ctx->row = (float *)malloc((size_t)d_model * sizeof(float)); + ctx->target = (float *)malloc((size_t)d_model * sizeof(float)); + ctx->expected = (float *)calloc((size_t)d_model, sizeof(float)); + if (!ctx->reader || !ctx->row || !ctx->target || !ctx->expected) { + orion_qwen_streaming_ce_context_close(ctx); + return NULL; + } + return ctx; +} + +float orion_qwen_streaming_ce_tied_embedding_ctx(OrionQwenStreamingCEContext *ctx, + const float *hidden, + int target_token, + float *d_hidden_out) { + if (!ctx || !ctx->reader || !hidden || !d_hidden_out || target_token < 0 || target_token >= ctx->vocab) { + return NAN; + } + + float max_logit = -INFINITY; + float sum_exp = 0.0f; + float target_logit = NAN; + int target_seen = 0; + memset(ctx->expected, 0, (size_t)ctx->d_model * sizeof(float)); + for (int tok = 0; tok < ctx->vocab; tok++) { + if (!orion_blob_row_reader_read_f32(ctx->reader, tok, ctx->row)) { + return NAN; + } + float dot = cblas_sdot(ctx->d_model, hidden, 1, ctx->row, 1); + if (tok == target_token) { + memcpy(ctx->target, ctx->row, (size_t)ctx->d_model * sizeof(float)); + target_logit = dot; + target_seen = 1; + } + + if (dot > max_logit) { + float scale = isfinite(max_logit) ? expf(max_logit - dot) : 0.0f; + for (int i = 0; i < ctx->d_model; i++) { + ctx->expected[i] = (ctx->expected[i] * scale) + ctx->row[i]; + } + sum_exp = (sum_exp * scale) + 1.0f; + max_logit = dot; + } else { + float expv = expf(dot - max_logit); + sum_exp += expv; + for (int i = 0; i < ctx->d_model; i++) ctx->expected[i] += expv * ctx->row[i]; + } + } + + if (!target_seen || !isfinite(target_logit) || !(sum_exp > 0.0f)) { + return NAN; + } + + float inv_sum = 1.0f / sum_exp; + for (int i = 0; i < ctx->d_model; i++) { + d_hidden_out[i] = ctx->expected[i] * inv_sum - ctx->target[i]; + } + + float loss = -(target_logit - max_logit - logf(sum_exp)); + return loss; +} + +void orion_qwen_streaming_ce_context_close(OrionQwenStreamingCEContext *ctx) { + if (!ctx) return; + if (ctx->reader) orion_blob_row_reader_close(ctx->reader); + free(ctx->row); + free(ctx->target); + free(ctx->expected); + free(ctx); +} diff --git a/kernels/training/qwen_lora_train.h b/kernels/training/qwen_lora_train.h new file mode 100644 index 0000000..bc6561e --- /dev/null +++ b/kernels/training/qwen_lora_train.h @@ -0,0 +1,127 @@ +#ifndef ORION_QWEN_LORA_TRAIN_H +#define ORION_QWEN_LORA_TRAIN_H + +#import +#import "../../model/weight_loader.h" +#import "qwen_lora_cpu_ops.h" + +typedef struct OrionQwen9BCPUTrainContext OrionQwen9BCPUTrainContext; + +typedef struct { + int step; + int layer_idx; + float lr; + float beta1; + float beta2; + float eps; + OrionLoRAMatrix q_proj; + OrionLoRAMatrix v_proj; + OrionQwenStreamingCEContext *ce_ctx; + int owns_ce_ctx; + OrionQwen9BCPUTrainContext *cpu_ctx; + int owns_cpu_ctx; +} OrionQwen9BLoRATrainer; + +typedef struct { + float loss; + double q_grad_abs_sum; + double v_grad_abs_sum; + double q_param_abs_sum; + double v_param_abs_sum; + int predicted_token; +} OrionQwen9BLoRASmokeResult; + +typedef struct { + int items_completed; + double loss_sum; + float loss_first; + float loss_last; + float loss_min; + float loss_max; + float loss_avg; + double q_grad_abs_sum_last; + double v_grad_abs_sum_last; + double q_param_abs_sum_last; + double v_param_abs_sum_last; + int predicted_token_last; +} OrionQwen9BLoRABatchResult; + +void orion_qwen9b_lora_trainer_init(OrionQwen9BLoRATrainer *trainer, + const OrionQwen35Manifest *manifest, + int layer_idx, + int rank, + float alpha, + float lr, + unsigned int seed); + +void orion_qwen9b_lora_trainer_free(OrionQwen9BLoRATrainer *trainer); + +int orion_qwen9b_lora_trainer_attach_ce_context(OrionQwen9BLoRATrainer *trainer, + const char *embed_blob_path, + const OrionQwen35Manifest *manifest); + +int orion_qwen9b_lora_trainer_attach_cpu_train_context(OrionQwen9BLoRATrainer *trainer, + const char *blob_dir, + const OrionQwen35Manifest *manifest); + +void orion_qwen9b_lora_trainer_zero_grad(OrionQwen9BLoRATrainer *trainer); + +void orion_qwen9b_lora_trainer_scale_grad(OrionQwen9BLoRATrainer *trainer, + float scale); + +void orion_qwen9b_lora_trainer_step(OrionQwen9BLoRATrainer *trainer); + +int orion_qwen9b_lora_trainer_save(const OrionQwen9BLoRATrainer *trainer, + const char *out_dir); + +int orion_qwen9b_lora_trainer_load(OrionQwen9BLoRATrainer *trainer, + const char *in_dir); + +int orion_qwen9b_lora_trainer_compare(const OrionQwen9BLoRATrainer *lhs, + const OrionQwen9BLoRATrainer *rhs, + float atol); + +int orion_qwen9b_lora_train_smoke1(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + OrionQwen9BLoRASmokeResult *out_result); + +int orion_qwen9b_lora_train_accumulate1(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + OrionQwen9BLoRASmokeResult *out_result); + +int orion_qwen9b_lora_train_accumulate_hidden1(const char *blob_dir, + const OrionQwen35Manifest *manifest, + OrionQwen9BLoRATrainer *trainer, + const float *hidden_in, + int target_token, + OrionQwen9BLoRASmokeResult *out_result); + +int orion_qwen9b_lora_train_hidden_batch(const char *blob_dir, + const OrionQwen35Manifest *manifest, + OrionQwen9BLoRATrainer *trainer, + const float *const *hidden_batch, + const int *target_tokens, + int item_count, + OrionQwen9BLoRABatchResult *out_result); + +int orion_qwen9b_lora_train_smoke1_ane_qv_base(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + OrionQwen9BLoRASmokeResult *out_result); + +int orion_qwen9b_lora_frozen_prefix_hidden(const char *blob_dir, + const OrionQwen35Manifest *manifest, + int input_token, + float *hidden_out); + +int orion_qwen9b_lora_ane_train_bridge_last_compile_cache_hit(void); +int orion_qwen9b_lora_ane_train_bridge_last_q_cache_hit(void); +int orion_qwen9b_lora_ane_train_bridge_last_kv_cache_hit(void); +const char *orion_qwen9b_lora_ane_train_bridge_last_compile_cache_source(void); + +#endif // ORION_QWEN_LORA_TRAIN_H diff --git a/kernels/training/qwen_lora_train.m b/kernels/training/qwen_lora_train.m new file mode 100644 index 0000000..592f659 --- /dev/null +++ b/kernels/training/qwen_lora_train.m @@ -0,0 +1,1430 @@ +#import "qwen_lora_train.h" +#import "../inference/qwen_cpu_ops.h" +#import "../training/stories_cpu_ops.h" +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#import +#import +#import +#import +#import +#import +#import + +static int file_exists(const char *path) { + return access(path, F_OK) == 0; +} + +static float *load_layer_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static int load_embed_row(const char *blob_dir, int token_id, int d_model, float *out_row) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + return orion_read_blob_row_f32(path, token_id, d_model, out_row); +} + +static void expand_grouped_value(const float *v, + int n_head, + int n_kv_head, + int head_dim, + float *attn_cat) { + int q_per_kv = n_head / n_kv_head; + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + memcpy(attn_cat + (size_t)h * head_dim, + v + (size_t)kv_head * head_dim, + (size_t)head_dim * sizeof(float)); + } +} + +static void reduce_grouped_value_grad(const float *d_attn_cat, + int n_head, + int n_kv_head, + int head_dim, + float *d_v) { + int q_per_kv = n_head / n_kv_head; + memset(d_v, 0, (size_t)n_kv_head * head_dim * sizeof(float)); + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + float *dst = d_v + (size_t)kv_head * head_dim; + const float *src = d_attn_cat + (size_t)h * head_dim; + for (int i = 0; i < head_dim; i++) dst[i] += src[i]; + } +} + +static int parse_index_list_local(const char *csv, int *out, int max_count, int limit) { + if (!csv || !*csv || !out || max_count <= 0) return 0; + char *copy = strdup(csv); + if (!copy) return 0; + int count = 0; + char *save = NULL; + for (char *tok = strtok_r(copy, ",", &save); tok && count < max_count; tok = strtok_r(NULL, ",", &save)) { + while (*tok && isspace((unsigned char)*tok)) tok++; + if (!*tok) continue; + char *end = NULL; + long value = strtol(tok, &end, 10); + if (end == tok || value < 0 || value >= limit) continue; + while (*end && isspace((unsigned char)*end)) end++; + if (*end != '\0') continue; + int duplicate = 0; + for (int i = 0; i < count; i++) { + if (out[i] == (int)value) { + duplicate = 1; + break; + } + } + if (!duplicate) out[count++] = (int)value; + } + free(copy); + return count; +} + +static int use_cpu_v_proj_override(void) { + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_proj_override(void) { + const char *source = getenv("ORION_Q_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_query_override(void) { + const char *source = getenv("ORION_Q_QUERY_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_gate_override(void) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int load_q_gate_cpu_channel_preset(const char *source, int *out_channels, int max_channels, int q_dim) { + static const int kSeed2[] = {143, 3994}; + static const int kSeed4[] = {143, 3994, 3768, 1353}; + const int *preset = NULL; + int preset_count = 0; + if (!source || !out_channels || max_channels <= 0) return 0; + if (strcmp(source, "seed2") == 0) { + preset = kSeed2; + preset_count = (int)(sizeof(kSeed2) / sizeof(kSeed2[0])); + } else if (strcmp(source, "seed4") == 0) { + preset = kSeed4; + preset_count = (int)(sizeof(kSeed4) / sizeof(kSeed4[0])); + } + if (!preset) return 0; + int count = 0; + for (int i = 0; i < preset_count && count < max_channels; i++) { + int channel = preset[i]; + if (channel < 0 || channel >= q_dim) continue; + out_channels[count++] = channel; + } + return count; +} + +static int load_q_gate_cpu_channel_override(int *out_channels, int max_channels, int q_dim) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + int preset_count = load_q_gate_cpu_channel_preset(source, out_channels, max_channels, q_dim); + if (preset_count > 0) return preset_count; + const char *csv = getenv("ORION_Q_GATE_CPU_CHANNELS"); + if (!csv || !*csv) return 0; + return parse_index_list_local(csv, out_channels, max_channels, q_dim); +} + +struct OrionQwen9BCPUTrainContext { + int d_model; + int d_ff; + int q_dim; + int kv_dim; + float *input_ln; + float *post_ln; + float *q_proj; + float *v_proj; + float *o_proj; + float *gate_proj; + float *up_proj; + float *down_proj; + float *final_norm_weight; + float *normed; + float *q_full; + float *v_raw; + float *attn_cat; + float *gated; + float *mixer; + float *hidden_mid; + float *post_norm; + float *mlp_out; + float *hidden_out; + float *final_norm; + float *last_hidden; + float *d_last_hidden; + float *d_hidden_out; + float *d_post_norm; + float *d_hidden_mid; + float *d_mixer; + float *d_gated; + float *d_attn_cat; + float *d_gate_half; + float *d_q_full; + float *d_v; + float *d_normed_from_q; + float *d_normed_from_v; + float *d_normed; + float *d_hidden_mid_from_post; + float *d_post_ln_weight_grad; + float *throwaway_weight_grad; + float *throwaway_dx; +}; + +static void orion_qwen9b_cpu_train_context_close(OrionQwen9BCPUTrainContext *ctx) { + if (!ctx) return; + free(ctx->input_ln); free(ctx->post_ln); free(ctx->q_proj); free(ctx->v_proj); free(ctx->o_proj); + free(ctx->gate_proj); free(ctx->up_proj); free(ctx->down_proj); free(ctx->final_norm_weight); + free(ctx->normed); free(ctx->q_full); free(ctx->v_raw); free(ctx->attn_cat); free(ctx->gated); + free(ctx->mixer); free(ctx->hidden_mid); free(ctx->post_norm); free(ctx->mlp_out); free(ctx->hidden_out); + free(ctx->final_norm); free(ctx->last_hidden); free(ctx->d_last_hidden); free(ctx->d_hidden_out); + free(ctx->d_post_norm); free(ctx->d_hidden_mid); free(ctx->d_mixer); free(ctx->d_gated); free(ctx->d_attn_cat); + free(ctx->d_gate_half); free(ctx->d_q_full); free(ctx->d_v); free(ctx->d_normed_from_q); free(ctx->d_normed_from_v); + free(ctx->d_normed); free(ctx->d_hidden_mid_from_post); free(ctx->d_post_ln_weight_grad); + free(ctx->throwaway_weight_grad); free(ctx->throwaway_dx); + free(ctx); +} + +static OrionQwen9BCPUTrainContext *orion_qwen9b_cpu_train_context_open(const char *blob_dir, + const OrionQwen35Manifest *manifest, + int layer_idx) { + if (!blob_dir || !manifest) return NULL; + + OrionQwen9BCPUTrainContext *ctx = (OrionQwen9BCPUTrainContext *)calloc(1, sizeof(OrionQwen9BCPUTrainContext)); + if (!ctx) return NULL; + + ctx->d_model = manifest->d_model; + ctx->d_ff = manifest->d_ff; + ctx->q_dim = manifest->n_head * manifest->head_dim; + ctx->kv_dim = manifest->n_kv_head * manifest->head_dim; + + ctx->input_ln = load_layer_exact(blob_dir, layer_idx, "input_layernorm.bin", ctx->d_model); + ctx->post_ln = load_layer_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", ctx->d_model); + ctx->q_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (ctx->q_dim * 2) * ctx->d_model); + ctx->v_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", ctx->kv_dim * ctx->d_model); + ctx->o_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", ctx->d_model * ctx->q_dim); + ctx->gate_proj = load_layer_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", ctx->d_ff * ctx->d_model); + ctx->up_proj = load_layer_exact(blob_dir, layer_idx, "mlp_up_proj.bin", ctx->d_ff * ctx->d_model); + ctx->down_proj = load_layer_exact(blob_dir, layer_idx, "mlp_down_proj.bin", ctx->d_model * ctx->d_ff); + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + ctx->final_norm_weight = orion_read_blob_f32_exact(final_norm_path, ctx->d_model); + + ctx->normed = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->q_full = (float *)calloc((size_t)(ctx->q_dim * 2), sizeof(float)); + ctx->v_raw = (float *)calloc((size_t)ctx->kv_dim, sizeof(float)); + ctx->attn_cat = (float *)calloc((size_t)ctx->q_dim, sizeof(float)); + ctx->gated = (float *)calloc((size_t)ctx->q_dim, sizeof(float)); + ctx->mixer = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->hidden_mid = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->post_norm = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->mlp_out = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->hidden_out = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->final_norm = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->last_hidden = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_last_hidden = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_hidden_out = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_post_norm = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_hidden_mid = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_mixer = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_gated = (float *)calloc((size_t)ctx->q_dim, sizeof(float)); + ctx->d_attn_cat = (float *)calloc((size_t)ctx->q_dim, sizeof(float)); + ctx->d_gate_half = (float *)calloc((size_t)ctx->q_dim, sizeof(float)); + ctx->d_q_full = (float *)calloc((size_t)(ctx->q_dim * 2), sizeof(float)); + ctx->d_v = (float *)calloc((size_t)ctx->kv_dim, sizeof(float)); + ctx->d_normed_from_q = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_normed_from_v = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_normed = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_hidden_mid_from_post = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->d_post_ln_weight_grad = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->throwaway_weight_grad = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + ctx->throwaway_dx = (float *)calloc((size_t)ctx->d_model, sizeof(float)); + + if (!ctx->input_ln || !ctx->post_ln || !ctx->q_proj || !ctx->v_proj || !ctx->o_proj || + !ctx->gate_proj || !ctx->up_proj || !ctx->down_proj || !ctx->final_norm_weight || + !ctx->normed || !ctx->q_full || !ctx->v_raw || !ctx->attn_cat || !ctx->gated || + !ctx->mixer || !ctx->hidden_mid || !ctx->post_norm || !ctx->mlp_out || !ctx->hidden_out || + !ctx->final_norm || !ctx->last_hidden || !ctx->d_last_hidden || !ctx->d_hidden_out || + !ctx->d_post_norm || !ctx->d_hidden_mid || !ctx->d_mixer || !ctx->d_gated || + !ctx->d_attn_cat || !ctx->d_gate_half || !ctx->d_q_full || !ctx->d_v || + !ctx->d_normed_from_q || !ctx->d_normed_from_v || !ctx->d_normed || + !ctx->d_hidden_mid_from_post || !ctx->d_post_ln_weight_grad || + !ctx->throwaway_weight_grad || !ctx->throwaway_dx) { + orion_qwen9b_cpu_train_context_close(ctx); + return NULL; + } + + return ctx; +} + +static void orion_qwen9b_cpu_train_context_reset(OrionQwen9BCPUTrainContext *ctx) { + if (!ctx) return; + memset(ctx->normed, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->q_full, 0, (size_t)(ctx->q_dim * 2) * sizeof(float)); + memset(ctx->v_raw, 0, (size_t)ctx->kv_dim * sizeof(float)); + memset(ctx->attn_cat, 0, (size_t)ctx->q_dim * sizeof(float)); + memset(ctx->gated, 0, (size_t)ctx->q_dim * sizeof(float)); + memset(ctx->mixer, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->hidden_mid, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->post_norm, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->mlp_out, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->hidden_out, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->final_norm, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->last_hidden, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_last_hidden, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_hidden_out, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_post_norm, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_hidden_mid, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_mixer, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_gated, 0, (size_t)ctx->q_dim * sizeof(float)); + memset(ctx->d_attn_cat, 0, (size_t)ctx->q_dim * sizeof(float)); + memset(ctx->d_gate_half, 0, (size_t)ctx->q_dim * sizeof(float)); + memset(ctx->d_q_full, 0, (size_t)(ctx->q_dim * 2) * sizeof(float)); + memset(ctx->d_v, 0, (size_t)ctx->kv_dim * sizeof(float)); + memset(ctx->d_normed_from_q, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_normed_from_v, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_normed, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_hidden_mid_from_post, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->d_post_ln_weight_grad, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->throwaway_weight_grad, 0, (size_t)ctx->d_model * sizeof(float)); + memset(ctx->throwaway_dx, 0, (size_t)ctx->d_model * sizeof(float)); +} + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_qproj_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +typedef struct { + OrionProgram *prog_q; + OrionProgram *prog_kv; + int layer_idx; + int bucket; + int d_model; + int q_dim; + int kv_dim; + int qkv_input_mode; + int q_uses_cpu_rms; + int kv_uses_cpu_rms; + int compile_cache_hit; + int q_cache_hit; + int kv_cache_hit; + int runtime_artifact_exported; + char compile_cache_source[64]; + char blob_dir[2048]; +} OrionQwenAneTrainBridge; + +static OrionQwenAneTrainBridge g_ane_train_bridge = {0}; + +enum { + ORION_QKV_INPUT_MODE_ANE_RMS = 0, + ORION_QKV_INPUT_MODE_CPU_RMS = 1, + ORION_QKV_INPUT_MODE_KV_CPU_RMS = 2, +}; + +static int parse_qkv_input_mode(void) { + const char *mode = getenv("ORION_QKV_INPUT_MODE"); + if (!mode || !*mode) return ORION_QKV_INPUT_MODE_ANE_RMS; + if (strcmp(mode, "cpu_rms") == 0) return ORION_QKV_INPUT_MODE_CPU_RMS; + if (strcmp(mode, "kv_cpu_rms") == 0) return ORION_QKV_INPUT_MODE_KV_CPU_RMS; + return ORION_QKV_INPUT_MODE_ANE_RMS; +} + +static int q_uses_cpu_rms_mode(int qkv_input_mode) { + return qkv_input_mode == ORION_QKV_INPUT_MODE_CPU_RMS; +} + +static int kv_uses_cpu_rms_mode(int qkv_input_mode) { + return qkv_input_mode == ORION_QKV_INPUT_MODE_CPU_RMS || + qkv_input_mode == ORION_QKV_INPUT_MODE_KV_CPU_RMS; +} + +static const char *compile_cache_root_env(void) { + const char *root = getenv("ORION_COMPILE_CACHE_DIR"); + if (root && *root) return root; + root = getenv("COMPILE_CACHE_DIR"); + return (root && *root) ? root : NULL; +} + +static const char *compile_cache_dataset_hash_env(void) { + const char *value = getenv("ORION_COMPILE_CACHE_DATASET_HASH"); + if (value && *value) return value; + value = getenv("DATASET_HASH"); + return (value && *value) ? value : "unknown"; +} + +static NSString *qkv_input_mode_name(int qkv_input_mode) { + switch (qkv_input_mode) { + case ORION_QKV_INPUT_MODE_CPU_RMS: return @"cpu_rms"; + case ORION_QKV_INPUT_MODE_KV_CPU_RMS: return @"kv_cpu_rms"; + default: return @"ane_rms"; + } +} + +static NSString *ane_train_bridge_compile_artifact_dir(const char *blob_dir, + int layer_idx, + int bucket, + int qkv_input_mode, + const char *component) { + const char *cache_root = compile_cache_root_env(); + if (!cache_root || !component || !*component) return nil; + NSString *blobBase = [[NSString stringWithUTF8String:blob_dir] lastPathComponent]; + NSString *datasetHash = [NSString stringWithUTF8String:compile_cache_dataset_hash_env()]; + NSString *modeName = qkv_input_mode_name(qkv_input_mode); + NSString *artifactName = [NSString stringWithFormat:@"qwen35_9b_lora_train_%@_%@_L%d_B%d_%@_%s", + blobBase, + datasetHash, + layer_idx, + bucket, + modeName, + component]; + return [[NSString stringWithUTF8String:cache_root] stringByAppendingPathComponent:artifactName]; +} + +static void ane_train_bridge_release(void) { + if (g_ane_train_bridge.prog_q) orion_release_program(g_ane_train_bridge.prog_q); + if (g_ane_train_bridge.prog_kv) orion_release_program(g_ane_train_bridge.prog_kv); + memset(&g_ane_train_bridge, 0, sizeof(g_ane_train_bridge)); +} + +static int ane_train_bridge_ensure(const char *blob_dir, + OrionQwen35Manifest *manifest, + int layer_idx, + int bucket) { + if (g_ane_train_bridge.prog_q && + g_ane_train_bridge.prog_kv && + g_ane_train_bridge.layer_idx == layer_idx && + g_ane_train_bridge.bucket == bucket && + g_ane_train_bridge.qkv_input_mode == parse_qkv_input_mode() && + strcmp(g_ane_train_bridge.blob_dir, blob_dir) == 0) { + return 1; + } + + ane_train_bridge_release(); + if (!orion_ane_init()) return 0; + + OrionModelConfig cfg = { + .n_layer = manifest->n_layer, + .n_head = manifest->n_head, + .n_kv_head = manifest->n_kv_head, + .d_model = manifest->d_model, + .head_dim = manifest->head_dim, + .hidden_dim = manifest->d_ff, + .vocab = manifest->vocab, + .max_seq = manifest->max_seq, + }; + NSString *blobDir = [NSString stringWithUTF8String:blob_dir]; + int qkv_input_mode = parse_qkv_input_mode(); + int q_uses_cpu_rms = q_uses_cpu_rms_mode(qkv_input_mode); + int kv_uses_cpu_rms = kv_uses_cpu_rms_mode(qkv_input_mode); + NSString *mil_q = compile_graph( + q_uses_cpu_rms + ? orion_frontend_qwen35_prefill_q_proj_linear_only(layer_idx, bucket, &cfg) + : orion_frontend_qwen35_prefill_q_proj(layer_idx, bucket, &cfg) + ); + NSString *mil_kv = compile_graph( + kv_uses_cpu_rms + ? orion_frontend_qwen35_prefill_kv_proj_linear_only(layer_idx, bucket, &cfg) + : orion_frontend_qwen35_prefill_kv_proj(layer_idx, bucket, &cfg) + ); + if (!mil_q || !mil_kv) return 0; + + NSDictionary *wdict_q = q_uses_cpu_rms ? build_qproj_linear_only_wdict(layer_idx, blobDir) : build_qproj_wdict(layer_idx, blobDir); + NSDictionary *wdict_kv = kv_uses_cpu_rms ? build_kv_linear_only_wdict(layer_idx, blobDir) : build_kv_wdict(layer_idx, blobDir); + NSString *artifact_q = ane_train_bridge_compile_artifact_dir(blob_dir, layer_idx, bucket, qkv_input_mode, "q"); + NSString *artifact_kv = ane_train_bridge_compile_artifact_dir(blob_dir, layer_idx, bucket, qkv_input_mode, "kv"); + int q_cache_hit = 0; + int kv_cache_hit = 0; + int runtime_artifact_exported = 0; + + if (artifact_q) { + g_ane_train_bridge.prog_q = orion_program_load_artifacts( + mil_q.UTF8String, + wdict_q, + artifact_q.UTF8String, + "qwen35_9b_lora_ane_train_q" + ); + q_cache_hit = g_ane_train_bridge.prog_q ? 1 : 0; + } + if (!g_ane_train_bridge.prog_q) { + g_ane_train_bridge.prog_q = orion_compile_mil( + mil_q.UTF8String, + wdict_q, + "qwen35_9b_lora_ane_train_q" + ); + if (g_ane_train_bridge.prog_q && artifact_q && + orion_program_export_artifacts(g_ane_train_bridge.prog_q, artifact_q.UTF8String)) { + runtime_artifact_exported = 1; + } + } + + if (artifact_kv) { + g_ane_train_bridge.prog_kv = orion_program_load_artifacts( + mil_kv.UTF8String, + wdict_kv, + artifact_kv.UTF8String, + "qwen35_9b_lora_ane_train_kv" + ); + kv_cache_hit = g_ane_train_bridge.prog_kv ? 1 : 0; + } + if (!g_ane_train_bridge.prog_kv) { + g_ane_train_bridge.prog_kv = orion_compile_mil( + mil_kv.UTF8String, + wdict_kv, + "qwen35_9b_lora_ane_train_kv" + ); + if (g_ane_train_bridge.prog_kv && artifact_kv && + orion_program_export_artifacts(g_ane_train_bridge.prog_kv, artifact_kv.UTF8String)) { + runtime_artifact_exported = 1; + } + } + if (!g_ane_train_bridge.prog_q || !g_ane_train_bridge.prog_kv) { + ane_train_bridge_release(); + return 0; + } + + g_ane_train_bridge.layer_idx = layer_idx; + g_ane_train_bridge.bucket = bucket; + g_ane_train_bridge.d_model = manifest->d_model; + g_ane_train_bridge.q_dim = manifest->n_head * manifest->head_dim; + g_ane_train_bridge.kv_dim = manifest->n_kv_head * manifest->head_dim; + g_ane_train_bridge.qkv_input_mode = qkv_input_mode; + g_ane_train_bridge.q_uses_cpu_rms = q_uses_cpu_rms; + g_ane_train_bridge.kv_uses_cpu_rms = kv_uses_cpu_rms; + g_ane_train_bridge.q_cache_hit = q_cache_hit; + g_ane_train_bridge.kv_cache_hit = kv_cache_hit; + g_ane_train_bridge.compile_cache_hit = q_cache_hit && kv_cache_hit; + g_ane_train_bridge.runtime_artifact_exported = runtime_artifact_exported; + if (g_ane_train_bridge.compile_cache_hit) { + strlcpy(g_ane_train_bridge.compile_cache_source, "runtime_artifact", sizeof(g_ane_train_bridge.compile_cache_source)); + } else if (q_cache_hit || kv_cache_hit) { + strlcpy(g_ane_train_bridge.compile_cache_source, "partial_runtime_artifact", sizeof(g_ane_train_bridge.compile_cache_source)); + } else { + strlcpy(g_ane_train_bridge.compile_cache_source, "fresh_compile", sizeof(g_ane_train_bridge.compile_cache_source)); + } + strlcpy(g_ane_train_bridge.blob_dir, blob_dir, sizeof(g_ane_train_bridge.blob_dir)); + return 1; +} + +int orion_qwen9b_lora_ane_train_bridge_last_compile_cache_hit(void) { + return g_ane_train_bridge.compile_cache_hit; +} + +int orion_qwen9b_lora_ane_train_bridge_last_q_cache_hit(void) { + return g_ane_train_bridge.q_cache_hit; +} + +int orion_qwen9b_lora_ane_train_bridge_last_kv_cache_hit(void) { + return g_ane_train_bridge.kv_cache_hit; +} + +const char *orion_qwen9b_lora_ane_train_bridge_last_compile_cache_source(void) { + return g_ane_train_bridge.compile_cache_source[0] ? g_ane_train_bridge.compile_cache_source : "unknown"; +} + +int orion_qwen9b_lora_frozen_prefix_hidden(const char *blob_dir, + const OrionQwen35Manifest *manifest, + int input_token, + float *hidden_out) { + const int seq_len = 1; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)d_model, sizeof(float)); + int ok = 0; + + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch) goto cleanup; + if (!load_embed_row(blob_dir, input_token, d_model, hidden)) goto cleanup; + + for (int layer_idx = 0; layer_idx < manifest->n_layer - 1; layer_idx++) { + float *input_ln = load_layer_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_layer_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_layer_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_layer_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_layer_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + + orion_qwen_cpu_rmsnorm(hidden, input_ln, d_model, 1e-6f, normed); + memset(mixer_out, 0, (size_t)d_model * sizeof(float)); + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + if (file_exists(full_q_path)) { + float *q_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_layer_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_layer_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_layer_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + char path_qkv[2048], path_out[2048], path_dt[2048], path_norm[2048], path_conv[2048]; + snprintf(path_qkv, sizeof(path_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + snprintf(path_out, sizeof(path_out), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx); + snprintf(path_dt, sizeof(path_dt), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + snprintf(path_norm, sizeof(path_norm), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx); + snprintf(path_conv, sizeof(path_conv), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + + int qkv_rows = orion_blob_element_count(path_qkv) / d_model; + int value_dim = orion_blob_element_count(path_out) / d_model; + int num_v_heads = orion_blob_element_count(path_dt); + int head_v_dim = orion_blob_element_count(path_norm); + int key_dim = (qkv_rows - value_dim) / 2; + int num_k_heads = num_v_heads; + int head_k_dim = key_dim / num_k_heads; + int conv_kernel = orion_blob_element_count(path_conv) / qkv_rows; + + float *in_proj_qkv = load_layer_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_layer_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_layer_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_layer_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_layer_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_layer_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_layer_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_layer_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_layer_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < d_model; i++) hidden[i] += mixer_out[i]; + orion_qwen_cpu_rmsnorm(hidden, post_ln, d_model, 1e-6f, scratch); + orion_qwen_cpu_swiglu_ffn(scratch, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out); + for (int i = 0; i < d_model; i++) hidden[i] += mlp_out[i]; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + memcpy(hidden_out, hidden, (size_t)d_model * sizeof(float)); + ok = 1; + +cleanup: + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + return ok; +} + +void orion_qwen9b_lora_trainer_init(OrionQwen9BLoRATrainer *trainer, + const OrionQwen35Manifest *manifest, + int layer_idx, + int rank, + float alpha, + float lr, + unsigned int seed) { + memset(trainer, 0, sizeof(*trainer)); + trainer->layer_idx = layer_idx; + trainer->lr = lr; + trainer->beta1 = 0.9f; + trainer->beta2 = 0.999f; + trainer->eps = 1e-8f; + trainer->step = 0; + trainer->ce_ctx = NULL; + trainer->owns_ce_ctx = 0; + trainer->cpu_ctx = NULL; + trainer->owns_cpu_ctx = 0; + + int q_dim = manifest->n_head * manifest->head_dim; + int kv_dim = manifest->n_kv_head * manifest->head_dim; + orion_qwen_lora_matrix_init(&trainer->q_proj, manifest->d_model, q_dim * 2, rank, alpha, seed ^ 0x13579BDFu); + orion_qwen_lora_matrix_init(&trainer->v_proj, manifest->d_model, kv_dim, rank, alpha, seed ^ 0x2468ACE0u); +} + +void orion_qwen9b_lora_trainer_free(OrionQwen9BLoRATrainer *trainer) { + if (!trainer) return; + if (trainer->owns_ce_ctx && trainer->ce_ctx) { + orion_qwen_streaming_ce_context_close(trainer->ce_ctx); + } + if (trainer->owns_cpu_ctx && trainer->cpu_ctx) { + orion_qwen9b_cpu_train_context_close(trainer->cpu_ctx); + } + orion_qwen_lora_matrix_free(&trainer->q_proj); + orion_qwen_lora_matrix_free(&trainer->v_proj); + memset(trainer, 0, sizeof(*trainer)); +} + +int orion_qwen9b_lora_trainer_attach_ce_context(OrionQwen9BLoRATrainer *trainer, + const char *embed_blob_path, + const OrionQwen35Manifest *manifest) { + if (!trainer || !embed_blob_path || !manifest) return 0; + if (trainer->owns_ce_ctx && trainer->ce_ctx) { + orion_qwen_streaming_ce_context_close(trainer->ce_ctx); + } + trainer->ce_ctx = orion_qwen_streaming_ce_context_open(embed_blob_path, manifest->d_model, manifest->vocab); + trainer->owns_ce_ctx = trainer->ce_ctx ? 1 : 0; + return trainer->ce_ctx != NULL; +} + +int orion_qwen9b_lora_trainer_attach_cpu_train_context(OrionQwen9BLoRATrainer *trainer, + const char *blob_dir, + const OrionQwen35Manifest *manifest) { + if (!trainer || !blob_dir || !manifest) return 0; + if (trainer->owns_cpu_ctx && trainer->cpu_ctx) { + orion_qwen9b_cpu_train_context_close(trainer->cpu_ctx); + } + trainer->cpu_ctx = orion_qwen9b_cpu_train_context_open(blob_dir, manifest, trainer->layer_idx); + trainer->owns_cpu_ctx = trainer->cpu_ctx ? 1 : 0; + return trainer->cpu_ctx != NULL; +} + +void orion_qwen9b_lora_trainer_zero_grad(OrionQwen9BLoRATrainer *trainer) { + orion_qwen_lora_matrix_zero_grad(&trainer->q_proj); + orion_qwen_lora_matrix_zero_grad(&trainer->v_proj); +} + +void orion_qwen9b_lora_trainer_scale_grad(OrionQwen9BLoRATrainer *trainer, + float scale) { + orion_qwen_lora_matrix_scale_grad(&trainer->q_proj, scale); + orion_qwen_lora_matrix_scale_grad(&trainer->v_proj, scale); +} + +void orion_qwen9b_lora_trainer_step(OrionQwen9BLoRATrainer *trainer) { + trainer->step += 1; + orion_cpu_adam_step(trainer->q_proj.a, trainer->q_proj.da, trainer->q_proj.ma, trainer->q_proj.va, + trainer->q_proj.rank * trainer->q_proj.in_dim, trainer->lr, trainer->beta1, trainer->beta2, trainer->eps, trainer->step); + orion_cpu_adam_step(trainer->q_proj.b, trainer->q_proj.db, trainer->q_proj.mb, trainer->q_proj.vb, + trainer->q_proj.out_dim * trainer->q_proj.rank, trainer->lr, trainer->beta1, trainer->beta2, trainer->eps, trainer->step); + orion_cpu_adam_step(trainer->v_proj.a, trainer->v_proj.da, trainer->v_proj.ma, trainer->v_proj.va, + trainer->v_proj.rank * trainer->v_proj.in_dim, trainer->lr, trainer->beta1, trainer->beta2, trainer->eps, trainer->step); + orion_cpu_adam_step(trainer->v_proj.b, trainer->v_proj.db, trainer->v_proj.mb, trainer->v_proj.vb, + trainer->v_proj.out_dim * trainer->v_proj.rank, trainer->lr, trainer->beta1, trainer->beta2, trainer->eps, trainer->step); + orion_qwen9b_lora_trainer_zero_grad(trainer); +} + +static int ensure_dir(const char *path) { + NSString *dir = [NSString stringWithUTF8String:path]; + NSError *error = nil; + return [[NSFileManager defaultManager] createDirectoryAtPath:dir + withIntermediateDirectories:YES + attributes:nil + error:&error]; +} + +int orion_qwen9b_lora_trainer_save(const OrionQwen9BLoRATrainer *trainer, + const char *out_dir) { + if (!ensure_dir(out_dir)) return 0; + + NSString *dir = [NSString stringWithUTF8String:out_dir]; + NSString *configPath = [dir stringByAppendingPathComponent:@"adapter_config.json"]; + NSString *statePath = [dir stringByAppendingPathComponent:@"trainer_state.json"]; + NSString *weightPath = [dir stringByAppendingPathComponent:@"adapter_weights.bin"]; + + NSDictionary *config = @{ + @"model_name": @"CRPG_Q3.5", + @"base_model": @"Qwen3.5-9B", + @"targets": @[@"q_proj", @"v_proj"], + @"rank": @(trainer->q_proj.rank), + @"alpha": @(trainer->q_proj.alpha), + @"dropout": @0.0, + @"layer_idx": @(trainer->layer_idx) + }; + NSDictionary *state = @{ + @"step": @(trainer->step), + @"lr": @(trainer->lr), + @"beta1": @(trainer->beta1), + @"beta2": @(trainer->beta2), + @"eps": @(trainer->eps) + }; + NSData *configData = [NSJSONSerialization dataWithJSONObject:config options:NSJSONWritingPrettyPrinted error:nil]; + NSData *stateData = [NSJSONSerialization dataWithJSONObject:state options:NSJSONWritingPrettyPrinted error:nil]; + if (![configData writeToFile:configPath atomically:YES]) return 0; + if (![stateData writeToFile:statePath atomically:YES]) return 0; + + FILE *f = fopen(weightPath.UTF8String, "wb"); + if (!f) return 0; + fwrite(trainer->q_proj.a, sizeof(float), (size_t)trainer->q_proj.rank * trainer->q_proj.in_dim, f); + fwrite(trainer->q_proj.b, sizeof(float), (size_t)trainer->q_proj.out_dim * trainer->q_proj.rank, f); + fwrite(trainer->v_proj.a, sizeof(float), (size_t)trainer->v_proj.rank * trainer->v_proj.in_dim, f); + fwrite(trainer->v_proj.b, sizeof(float), (size_t)trainer->v_proj.out_dim * trainer->v_proj.rank, f); + fclose(f); + return 1; +} + +int orion_qwen9b_lora_trainer_load(OrionQwen9BLoRATrainer *trainer, + const char *in_dir) { + NSString *dir = [NSString stringWithUTF8String:in_dir]; + NSString *weightPath = [dir stringByAppendingPathComponent:@"adapter_weights.bin"]; + NSString *statePath = [dir stringByAppendingPathComponent:@"trainer_state.json"]; + FILE *f = fopen(weightPath.UTF8String, "rb"); + if (!f) return 0; + size_t q_a = (size_t)trainer->q_proj.rank * trainer->q_proj.in_dim; + size_t q_b = (size_t)trainer->q_proj.out_dim * trainer->q_proj.rank; + size_t v_a = (size_t)trainer->v_proj.rank * trainer->v_proj.in_dim; + size_t v_b = (size_t)trainer->v_proj.out_dim * trainer->v_proj.rank; + int ok = fread(trainer->q_proj.a, sizeof(float), q_a, f) == q_a && + fread(trainer->q_proj.b, sizeof(float), q_b, f) == q_b && + fread(trainer->v_proj.a, sizeof(float), v_a, f) == v_a && + fread(trainer->v_proj.b, sizeof(float), v_b, f) == v_b; + fclose(f); + if (!ok) return 0; + + NSData *stateData = [NSData dataWithContentsOfFile:statePath]; + if (!stateData) return 0; + NSDictionary *state = [NSJSONSerialization JSONObjectWithData:stateData options:0 error:nil]; + trainer->step = [state[@"step"] intValue]; + trainer->lr = [state[@"lr"] floatValue]; + trainer->beta1 = [state[@"beta1"] floatValue]; + trainer->beta2 = [state[@"beta2"] floatValue]; + trainer->eps = [state[@"eps"] floatValue]; + return 1; +} + +int orion_qwen9b_lora_trainer_compare(const OrionQwen9BLoRATrainer *lhs, + const OrionQwen9BLoRATrainer *rhs, + float atol) { + if (lhs->layer_idx != rhs->layer_idx || lhs->step != rhs->step) return 0; + size_t q_a = (size_t)lhs->q_proj.rank * lhs->q_proj.in_dim; + size_t q_b = (size_t)lhs->q_proj.out_dim * lhs->q_proj.rank; + size_t v_a = (size_t)lhs->v_proj.rank * lhs->v_proj.in_dim; + size_t v_b = (size_t)lhs->v_proj.out_dim * lhs->v_proj.rank; + const float *pairs[] = { lhs->q_proj.a, rhs->q_proj.a, lhs->q_proj.b, rhs->q_proj.b, + lhs->v_proj.a, rhs->v_proj.a, lhs->v_proj.b, rhs->v_proj.b }; + const size_t counts[] = { q_a, q_b, v_a, v_b }; + for (int block = 0; block < 4; block++) { + const float *a = pairs[block * 2 + 0]; + const float *b = pairs[block * 2 + 1]; + for (size_t i = 0; i < counts[block]; i++) { + if (fabsf(a[i] - b[i]) > atol) return 0; + } + } + return 1; +} + +static int orion_qwen9b_lora_train_with_hidden_internal(const char *blob_dir, + const OrionQwen35Manifest *manifest, + OrionQwen9BLoRATrainer *trainer, + const float *hidden_in, + int target_token, + int apply_step, + OrionQwen9BLoRASmokeResult *out_result) { + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + OrionQwen9BCPUTrainContext *cpu_ctx = trainer->cpu_ctx; + + if (cpu_ctx) { + orion_qwen9b_cpu_train_context_reset(cpu_ctx); + } + + float *normed = cpu_ctx ? cpu_ctx->normed : (float *)calloc((size_t)d_model, sizeof(float)); + float *q_full = cpu_ctx ? cpu_ctx->q_full : (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *v_raw = cpu_ctx ? cpu_ctx->v_raw : (float *)calloc((size_t)kv_dim, sizeof(float)); + float *attn_cat = cpu_ctx ? cpu_ctx->attn_cat : (float *)calloc((size_t)q_dim, sizeof(float)); + float *gated = cpu_ctx ? cpu_ctx->gated : (float *)calloc((size_t)q_dim, sizeof(float)); + float *mixer = cpu_ctx ? cpu_ctx->mixer : (float *)calloc((size_t)d_model, sizeof(float)); + float *hidden_mid = cpu_ctx ? cpu_ctx->hidden_mid : (float *)calloc((size_t)d_model, sizeof(float)); + float *post_norm = cpu_ctx ? cpu_ctx->post_norm : (float *)calloc((size_t)d_model, sizeof(float)); + float *mlp_out = cpu_ctx ? cpu_ctx->mlp_out : (float *)calloc((size_t)d_model, sizeof(float)); + float *hidden_out = cpu_ctx ? cpu_ctx->hidden_out : (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = cpu_ctx ? cpu_ctx->final_norm : (float *)calloc((size_t)d_model, sizeof(float)); + float *last_hidden = cpu_ctx ? cpu_ctx->last_hidden : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_last_hidden = cpu_ctx ? cpu_ctx->d_last_hidden : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_out = cpu_ctx ? cpu_ctx->d_hidden_out : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_post_norm = cpu_ctx ? cpu_ctx->d_post_norm : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_mid = cpu_ctx ? cpu_ctx->d_hidden_mid : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_mixer = cpu_ctx ? cpu_ctx->d_mixer : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_gated = cpu_ctx ? cpu_ctx->d_gated : (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_attn_cat = cpu_ctx ? cpu_ctx->d_attn_cat : (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_gate_half = cpu_ctx ? cpu_ctx->d_gate_half : (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_q_full = cpu_ctx ? cpu_ctx->d_q_full : (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *d_v = cpu_ctx ? cpu_ctx->d_v : (float *)calloc((size_t)kv_dim, sizeof(float)); + float *d_normed_from_q = cpu_ctx ? cpu_ctx->d_normed_from_q : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_normed_from_v = cpu_ctx ? cpu_ctx->d_normed_from_v : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_normed = cpu_ctx ? cpu_ctx->d_normed : (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_mid_from_post = cpu_ctx ? cpu_ctx->d_hidden_mid_from_post : NULL; + float *d_post_ln_weight_grad = cpu_ctx ? cpu_ctx->d_post_ln_weight_grad : NULL; + + float *input_ln = cpu_ctx ? cpu_ctx->input_ln : load_layer_exact(blob_dir, trainer->layer_idx, "input_layernorm.bin", d_model); + float *post_ln = cpu_ctx ? cpu_ctx->post_ln : load_layer_exact(blob_dir, trainer->layer_idx, "post_attention_layernorm.bin", d_model); + float *q_proj = cpu_ctx ? cpu_ctx->q_proj : load_layer_exact(blob_dir, trainer->layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *v_proj = cpu_ctx ? cpu_ctx->v_proj : load_layer_exact(blob_dir, trainer->layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = cpu_ctx ? cpu_ctx->o_proj : load_layer_exact(blob_dir, trainer->layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *gate_proj = cpu_ctx ? cpu_ctx->gate_proj : load_layer_exact(blob_dir, trainer->layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = cpu_ctx ? cpu_ctx->up_proj : load_layer_exact(blob_dir, trainer->layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = cpu_ctx ? cpu_ctx->down_proj : load_layer_exact(blob_dir, trainer->layer_idx, "mlp_down_proj.bin", d_model * d_ff); + + char final_norm_path[2048], embed_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + float *final_norm_weight = cpu_ctx ? cpu_ctx->final_norm_weight : orion_read_blob_f32_exact(final_norm_path, d_model); + + if (!cpu_ctx && (!normed || !q_full || !v_raw || !attn_cat || !gated || !mixer || !hidden_mid || + !post_norm || !mlp_out || !hidden_out || !final_norm || !last_hidden || !d_last_hidden || + !d_hidden_out || !d_post_norm || !d_hidden_mid || !d_mixer || !d_gated || !d_attn_cat || + !d_gate_half || !d_q_full || !d_v || !d_normed_from_q || !d_normed_from_v || !d_normed || + !input_ln || !post_ln || !q_proj || !v_proj || !o_proj || !gate_proj || !up_proj || + !down_proj || !final_norm_weight)) goto fail; + + orion_qwen_cpu_rmsnorm(hidden_in, input_ln, d_model, 1e-6f, normed); + + orion_qwen_lora_linear_forward(normed, q_proj, &trainer->q_proj, q_full); + orion_qwen_lora_linear_forward(normed, v_proj, &trainer->v_proj, v_raw); + + expand_grouped_value(v_raw, n_head, n_kv_head, head_dim, attn_cat); + for (int i = 0; i < q_dim; i++) { + float gate = 1.0f / (1.0f + expf(-q_full[q_dim + i])); + gated[i] = attn_cat[i] * gate; + } + cblas_sgemv(CblasRowMajor, CblasTrans, + d_model, q_dim, + 1.0f, o_proj, q_dim, gated, 1, + 0.0f, mixer, 1); + + for (int i = 0; i < d_model; i++) hidden_mid[i] = hidden_in[i] + mixer[i]; + orion_qwen_cpu_rmsnorm(hidden_mid, post_ln, d_model, 1e-6f, post_norm); + orion_qwen_cpu_swiglu_ffn(post_norm, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out); + for (int i = 0; i < d_model; i++) hidden_out[i] = hidden_mid[i] + mlp_out[i]; + orion_qwen_cpu_rmsnorm(hidden_out, final_norm_weight, d_model, 1e-6f, last_hidden); + + float loss = trainer->ce_ctx + ? orion_qwen_streaming_ce_tied_embedding_ctx(trainer->ce_ctx, last_hidden, target_token, d_last_hidden) + : orion_qwen_cpu_streaming_ce_tied_embedding(embed_path, last_hidden, d_model, manifest->vocab, target_token, d_last_hidden); + if (!isfinite(loss)) goto fail; + + orion_cpu_rmsnorm_bwd(d_hidden_out, final_norm, d_last_hidden, hidden_out, final_norm_weight, d_model, 1, 1e-6f); + + orion_qwen_cpu_swiglu_ffn_bwd(post_norm, gate_proj, up_proj, down_proj, d_model, d_ff, d_hidden_out, d_post_norm); + if (!cpu_ctx) { + d_hidden_mid_from_post = (float *)calloc((size_t)d_model, sizeof(float)); + d_post_ln_weight_grad = (float *)calloc((size_t)d_model, sizeof(float)); + } + if (!d_hidden_mid_from_post || !d_post_ln_weight_grad) goto fail; + orion_cpu_rmsnorm_bwd(d_hidden_mid_from_post, d_post_ln_weight_grad, d_post_norm, hidden_mid, post_ln, d_model, 1, 1e-6f); + for (int i = 0; i < d_model; i++) d_hidden_mid[i] = d_hidden_out[i] + d_hidden_mid_from_post[i]; + + memcpy(d_mixer, d_hidden_mid, (size_t)d_model * sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + d_model, q_dim, + 1.0f, o_proj, q_dim, d_mixer, 1, + 0.0f, d_gated, 1); + for (int i = 0; i < q_dim; i++) { + float gate_pre = q_full[q_dim + i]; + float gate = 1.0f / (1.0f + expf(-gate_pre)); + d_attn_cat[i] = d_gated[i] * gate; + d_gate_half[i] = d_gated[i] * attn_cat[i] * gate * (1.0f - gate); + d_q_full[q_dim + i] = d_gate_half[i]; + } + reduce_grouped_value_grad(d_attn_cat, n_head, n_kv_head, head_dim, d_v); + + orion_qwen_lora_linear_backward(normed, q_proj, &trainer->q_proj, d_q_full, d_normed_from_q); + orion_qwen_lora_linear_backward(normed, v_proj, &trainer->v_proj, d_v, d_normed_from_v); + for (int i = 0; i < d_model; i++) d_normed[i] = d_normed_from_q[i] + d_normed_from_v[i]; + + { + float *throwaway_weight_grad = cpu_ctx ? cpu_ctx->throwaway_weight_grad : (float *)calloc((size_t)d_model, sizeof(float)); + float *throwaway_dx = cpu_ctx ? cpu_ctx->throwaway_dx : (float *)calloc((size_t)d_model, sizeof(float)); + if (!throwaway_weight_grad || !throwaway_dx) { + if (!cpu_ctx) { + free(throwaway_weight_grad); + free(throwaway_dx); + } + goto fail; + } + orion_cpu_rmsnorm_bwd(throwaway_dx, throwaway_weight_grad, d_normed, hidden_in, input_ln, d_model, 1, 1e-6f); + if (!cpu_ctx) { + free(throwaway_weight_grad); + free(throwaway_dx); + } + } + + float q_grad_abs_sum = orion_qwen_lora_grad_abs_sum(&trainer->q_proj); + float v_grad_abs_sum = orion_qwen_lora_grad_abs_sum(&trainer->v_proj); + + if (apply_step) orion_qwen9b_lora_trainer_step(trainer); + + memset(out_result, 0, sizeof(*out_result)); + out_result->loss = loss; + out_result->q_grad_abs_sum = q_grad_abs_sum; + out_result->v_grad_abs_sum = v_grad_abs_sum; + out_result->q_param_abs_sum = orion_qwen_lora_abs_sum(trainer->q_proj.a, trainer->q_proj.rank * trainer->q_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer->q_proj.b, trainer->q_proj.out_dim * trainer->q_proj.rank); + out_result->v_param_abs_sum = orion_qwen_lora_abs_sum(trainer->v_proj.a, trainer->v_proj.rank * trainer->v_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer->v_proj.b, trainer->v_proj.out_dim * trainer->v_proj.rank); + out_result->predicted_token = -1; + + if (!cpu_ctx) { + free(normed); free(q_full); free(v_raw); free(attn_cat); free(gated); free(mixer); + free(hidden_mid); free(post_norm); free(mlp_out); free(hidden_out); free(final_norm); free(last_hidden); + free(d_last_hidden); free(d_hidden_out); free(d_post_norm); free(d_hidden_mid); free(d_mixer); + free(d_gated); free(d_attn_cat); free(d_gate_half); free(d_q_full); free(d_v); free(d_normed_from_q); + free(d_normed_from_v); free(d_normed); + free(d_hidden_mid_from_post); free(d_post_ln_weight_grad); + free(input_ln); free(post_ln); free(q_proj); free(v_proj); free(o_proj); free(gate_proj); free(up_proj); free(down_proj); + free(final_norm_weight); + } + return 1; + +fail: + if (!cpu_ctx) { + free(normed); free(q_full); free(v_raw); free(attn_cat); free(gated); free(mixer); + free(hidden_mid); free(post_norm); free(mlp_out); free(hidden_out); free(final_norm); free(last_hidden); + free(d_last_hidden); free(d_hidden_out); free(d_post_norm); free(d_hidden_mid); free(d_mixer); + free(d_gated); free(d_attn_cat); free(d_gate_half); free(d_q_full); free(d_v); free(d_normed_from_q); + free(d_normed_from_v); free(d_normed); + free(d_hidden_mid_from_post); free(d_post_ln_weight_grad); + free(input_ln); free(post_ln); free(q_proj); free(v_proj); free(o_proj); free(gate_proj); free(up_proj); free(down_proj); + free(final_norm_weight); + } + return 0; +} + +static int orion_qwen9b_lora_train_single_internal(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + int apply_step, + OrionQwen9BLoRASmokeResult *out_result) { + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) return 0; + + int ok = 0; + float *hidden_in = (float *)calloc((size_t)manifest->d_model, sizeof(float)); + if (!hidden_in) goto done; + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, input_token, hidden_in)) goto done; + ok = orion_qwen9b_lora_train_with_hidden_internal(blob_dir, manifest, trainer, hidden_in, target_token, apply_step, out_result); + +done: + free(hidden_in); + orion_qwen35_manifest_free(manifest); + return ok; +} + +int orion_qwen9b_lora_train_smoke1(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + OrionQwen9BLoRASmokeResult *out_result) { + return orion_qwen9b_lora_train_single_internal(blob_dir, trainer, input_token, target_token, 1, out_result); +} + +int orion_qwen9b_lora_train_accumulate1(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + OrionQwen9BLoRASmokeResult *out_result) { + return orion_qwen9b_lora_train_single_internal(blob_dir, trainer, input_token, target_token, 0, out_result); +} + +int orion_qwen9b_lora_train_accumulate_hidden1(const char *blob_dir, + const OrionQwen35Manifest *manifest, + OrionQwen9BLoRATrainer *trainer, + const float *hidden_in, + int target_token, + OrionQwen9BLoRASmokeResult *out_result) { + if (!manifest || !hidden_in) return 0; + return orion_qwen9b_lora_train_with_hidden_internal(blob_dir, manifest, trainer, hidden_in, target_token, 0, out_result); +} + +int orion_qwen9b_lora_train_hidden_batch(const char *blob_dir, + const OrionQwen35Manifest *manifest, + OrionQwen9BLoRATrainer *trainer, + const float *const *hidden_batch, + const int *target_tokens, + int item_count, + OrionQwen9BLoRABatchResult *out_result) { + if (!blob_dir || !manifest || !trainer || !hidden_batch || !target_tokens || !out_result || item_count <= 0) { + return 0; + } + + orion_qwen9b_lora_trainer_zero_grad(trainer); + + double loss_sum = 0.0; + float loss_first = NAN; + float loss_last = NAN; + float loss_min = INFINITY; + float loss_max = -INFINITY; + OrionQwen9BLoRASmokeResult step_result; + + for (int idx = 0; idx < item_count; idx++) { + const float *hidden_in = hidden_batch[idx]; + if (!hidden_in) return 0; + if (!orion_qwen9b_lora_train_with_hidden_internal(blob_dir, manifest, trainer, hidden_in, + target_tokens[idx], 0, &step_result)) { + return 0; + } + if (!isfinite(step_result.loss)) return 0; + if (idx == 0) loss_first = step_result.loss; + loss_last = step_result.loss; + if (step_result.loss < loss_min) loss_min = step_result.loss; + if (step_result.loss > loss_max) loss_max = step_result.loss; + loss_sum += step_result.loss; + } + + orion_qwen9b_lora_trainer_scale_grad(trainer, 1.0f / (float)item_count); + double q_grad_last = orion_qwen_lora_grad_abs_sum(&trainer->q_proj); + double v_grad_last = orion_qwen_lora_grad_abs_sum(&trainer->v_proj); + orion_qwen9b_lora_trainer_step(trainer); + double q_param_last = orion_qwen_lora_abs_sum(trainer->q_proj.a, trainer->q_proj.rank * trainer->q_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer->q_proj.b, trainer->q_proj.out_dim * trainer->q_proj.rank); + double v_param_last = orion_qwen_lora_abs_sum(trainer->v_proj.a, trainer->v_proj.rank * trainer->v_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer->v_proj.b, trainer->v_proj.out_dim * trainer->v_proj.rank); + + memset(out_result, 0, sizeof(*out_result)); + out_result->items_completed = item_count; + out_result->loss_sum = loss_sum; + out_result->loss_first = loss_first; + out_result->loss_last = loss_last; + out_result->loss_min = loss_min; + out_result->loss_max = loss_max; + out_result->loss_avg = (float)(loss_sum / (double)item_count); + out_result->q_grad_abs_sum_last = q_grad_last; + out_result->v_grad_abs_sum_last = v_grad_last; + out_result->q_param_abs_sum_last = q_param_last; + out_result->v_param_abs_sum_last = v_param_last; + out_result->predicted_token_last = step_result.predicted_token; + return 1; +} + +int orion_qwen9b_lora_train_smoke1_ane_qv_base(const char *blob_dir, + OrionQwen9BLoRATrainer *trainer, + int input_token, + int target_token, + OrionQwen9BLoRASmokeResult *out_result) { + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) return 0; + + const int bucket = 32; + const int seq_len = 1; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int q_proj_uses_cpu = use_cpu_q_proj_override(); + const int q_query_uses_cpu = use_cpu_q_query_override(); + const int q_gate_uses_cpu = use_cpu_q_gate_override(); + const int v_proj_uses_cpu = use_cpu_v_proj_override(); + int q_gate_cpu_channels[32] = {0}; + const int q_gate_cpu_channel_count = load_q_gate_cpu_channel_override(q_gate_cpu_channels, 32, q_dim); + + float *hidden_in = (float *)calloc((size_t)d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *base_q = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *base_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *cpu_q = (q_proj_uses_cpu || q_query_uses_cpu || q_gate_uses_cpu || q_gate_cpu_channel_count > 0) + ? (float *)calloc((size_t)(q_dim * 2), sizeof(float)) + : NULL; + float *cpu_v = v_proj_uses_cpu ? (float *)calloc((size_t)kv_dim, sizeof(float)) : NULL; + float *delta_q = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *delta_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *q_full = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *v_raw = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *attn_cat = (float *)calloc((size_t)q_dim, sizeof(float)); + float *gated = (float *)calloc((size_t)q_dim, sizeof(float)); + float *mixer = (float *)calloc((size_t)d_model, sizeof(float)); + float *hidden_mid = (float *)calloc((size_t)d_model, sizeof(float)); + float *post_norm = (float *)calloc((size_t)d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *hidden_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = (float *)calloc((size_t)d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_post_norm = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_mid = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_mixer = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_gated = (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_attn_cat = (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_gate_half = (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_q_full = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *d_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *d_normed_from_q = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_normed_from_v = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_mid_from_post = NULL; + float *d_post_ln_weight_grad = NULL; + float *input_ln = load_layer_exact(blob_dir, trainer->layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_layer_exact(blob_dir, trainer->layer_idx, "post_attention_layernorm.bin", d_model); + float *q_proj = load_layer_exact(blob_dir, trainer->layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *v_proj = load_layer_exact(blob_dir, trainer->layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_layer_exact(blob_dir, trainer->layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *gate_proj = load_layer_exact(blob_dir, trainer->layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_layer_exact(blob_dir, trainer->layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_layer_exact(blob_dir, trainer->layer_idx, "mlp_down_proj.bin", d_model * d_ff); + char final_norm_path[2048], embed_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + float *final_norm_weight = orion_read_blob_f32_exact(final_norm_path, d_model); + IOSurfaceRef ioIn = NULL; + IOSurfaceRef ioNormedIn = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + int ok = 0; + float loss = NAN; + + if (!hidden_in || !normed || !base_q || !base_v || !delta_q || !delta_v || !q_full || !v_raw || + !attn_cat || !gated || !mixer || !hidden_mid || !post_norm || !mlp_out || !hidden_out || + !final_norm || !last_hidden || !d_last_hidden || !d_hidden_out || !d_post_norm || + !d_hidden_mid || !d_mixer || !d_gated || !d_attn_cat || !d_gate_half || !d_q_full || + !d_v || !d_normed_from_q || !d_normed_from_v || !d_normed || !input_ln || !post_ln || + !q_proj || !v_proj || !o_proj || !gate_proj || !up_proj || !down_proj || !final_norm_weight || + ((q_proj_uses_cpu || q_query_uses_cpu || q_gate_uses_cpu || q_gate_cpu_channel_count > 0) && !cpu_q) || + (v_proj_uses_cpu && !cpu_v)) goto fail; + + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, input_token, hidden_in)) goto fail; + orion_qwen_cpu_rmsnorm(hidden_in, input_ln, d_model, 1e-6f, normed); + if (!ane_train_bridge_ensure(blob_dir, manifest, trainer->layer_idx, bucket)) goto fail; + + ioIn = make_cpu_seq_input_surface(hidden_in, seq_len, bucket, d_model); + ioNormedIn = make_cpu_seq_input_surface(normed, seq_len, bucket, d_model); + ioQ = make_f32_surface((q_dim * 2) * bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bucket, 0.0f); + if (!ioIn || !ioNormedIn || !ioQ || !ioK || !ioV) goto fail; + + IOSurfaceRef q_ins[] = { g_ane_train_bridge.q_uses_cpu_rms ? ioNormedIn : ioIn }; + IOSurfaceRef kv_ins[] = { g_ane_train_bridge.kv_uses_cpu_rms ? ioNormedIn : ioIn }; + IOSurfaceRef outsQ[] = { ioQ }; + IOSurfaceRef outsKV[] = { ioK, ioV }; + if (!orion_eval(g_ane_train_bridge.prog_q, q_ins, 1, outsQ, 1)) goto fail; + if (!orion_eval(g_ane_train_bridge.prog_kv, kv_ins, 1, outsKV, 2)) goto fail; + + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bucket, base_q); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bucket, base_v); + if (cpu_q) { + cblas_sgemv(CblasRowMajor, CblasNoTrans, + q_dim * 2, d_model, + 1.0f, q_proj, d_model, normed, 1, + 0.0f, cpu_q, 1); + if (q_proj_uses_cpu || q_query_uses_cpu) { + memcpy(base_q, cpu_q, (size_t)q_dim * sizeof(float)); + } + if (q_proj_uses_cpu || q_gate_uses_cpu) { + memcpy(base_q + q_dim, cpu_q + q_dim, (size_t)q_dim * sizeof(float)); + } else { + for (int i = 0; i < q_gate_cpu_channel_count; i++) { + int channel = q_gate_cpu_channels[i]; + base_q[q_dim + channel] = cpu_q[q_dim + channel]; + } + } + } + if (cpu_v) { + cblas_sgemv(CblasRowMajor, CblasNoTrans, + kv_dim, d_model, + 1.0f, v_proj, d_model, normed, 1, + 0.0f, cpu_v, 1); + memcpy(base_v, cpu_v, (size_t)kv_dim * sizeof(float)); + } + orion_qwen_lora_linear_delta_forward(normed, &trainer->q_proj, delta_q); + orion_qwen_lora_linear_delta_forward(normed, &trainer->v_proj, delta_v); + for (int i = 0; i < q_dim * 2; i++) q_full[i] = base_q[i] + delta_q[i]; + for (int i = 0; i < kv_dim; i++) v_raw[i] = base_v[i] + delta_v[i]; + + expand_grouped_value(v_raw, n_head, n_kv_head, head_dim, attn_cat); + for (int i = 0; i < q_dim; i++) { + float gate = 1.0f / (1.0f + expf(-q_full[q_dim + i])); + gated[i] = attn_cat[i] * gate; + } + cblas_sgemv(CblasRowMajor, CblasTrans, + d_model, q_dim, + 1.0f, o_proj, q_dim, gated, 1, + 0.0f, mixer, 1); + + for (int i = 0; i < d_model; i++) hidden_mid[i] = hidden_in[i] + mixer[i]; + orion_qwen_cpu_rmsnorm(hidden_mid, post_ln, d_model, 1e-6f, post_norm); + orion_qwen_cpu_swiglu_ffn(post_norm, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out); + for (int i = 0; i < d_model; i++) hidden_out[i] = hidden_mid[i] + mlp_out[i]; + orion_qwen_cpu_rmsnorm(hidden_out, final_norm_weight, d_model, 1e-6f, last_hidden); + + loss = trainer->ce_ctx + ? orion_qwen_streaming_ce_tied_embedding_ctx(trainer->ce_ctx, last_hidden, target_token, d_last_hidden) + : orion_qwen_cpu_streaming_ce_tied_embedding(embed_path, last_hidden, d_model, manifest->vocab, target_token, d_last_hidden); + if (!isfinite(loss)) goto fail; + + orion_cpu_rmsnorm_bwd(d_hidden_out, final_norm, d_last_hidden, hidden_out, final_norm_weight, d_model, 1, 1e-6f); + orion_qwen_cpu_swiglu_ffn_bwd(post_norm, gate_proj, up_proj, down_proj, d_model, d_ff, d_hidden_out, d_post_norm); + d_hidden_mid_from_post = (float *)calloc((size_t)d_model, sizeof(float)); + d_post_ln_weight_grad = (float *)calloc((size_t)d_model, sizeof(float)); + orion_cpu_rmsnorm_bwd(d_hidden_mid_from_post, d_post_ln_weight_grad, d_post_norm, hidden_mid, post_ln, d_model, 1, 1e-6f); + for (int i = 0; i < d_model; i++) d_hidden_mid[i] = d_hidden_out[i] + d_hidden_mid_from_post[i]; + + memcpy(d_mixer, d_hidden_mid, (size_t)d_model * sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + d_model, q_dim, + 1.0f, o_proj, q_dim, d_mixer, 1, + 0.0f, d_gated, 1); + for (int i = 0; i < q_dim; i++) { + float gate_pre = q_full[q_dim + i]; + float gate = 1.0f / (1.0f + expf(-gate_pre)); + d_attn_cat[i] = d_gated[i] * gate; + d_gate_half[i] = d_gated[i] * attn_cat[i] * gate * (1.0f - gate); + d_q_full[q_dim + i] = d_gate_half[i]; + } + reduce_grouped_value_grad(d_attn_cat, n_head, n_kv_head, head_dim, d_v); + + orion_qwen_lora_linear_backward(normed, q_proj, &trainer->q_proj, d_q_full, d_normed_from_q); + orion_qwen_lora_linear_backward(normed, v_proj, &trainer->v_proj, d_v, d_normed_from_v); + for (int i = 0; i < d_model; i++) d_normed[i] = d_normed_from_q[i] + d_normed_from_v[i]; + + { + float *throwaway_weight_grad = (float *)calloc((size_t)d_model, sizeof(float)); + float *throwaway_dx = (float *)calloc((size_t)d_model, sizeof(float)); + orion_cpu_rmsnorm_bwd(throwaway_dx, throwaway_weight_grad, d_normed, hidden_in, input_ln, d_model, 1, 1e-6f); + free(throwaway_weight_grad); + free(throwaway_dx); + } + + memset(out_result, 0, sizeof(*out_result)); + out_result->loss = loss; + out_result->q_grad_abs_sum = orion_qwen_lora_grad_abs_sum(&trainer->q_proj); + out_result->v_grad_abs_sum = orion_qwen_lora_grad_abs_sum(&trainer->v_proj); + orion_qwen9b_lora_trainer_step(trainer); + out_result->q_param_abs_sum = orion_qwen_lora_abs_sum(trainer->q_proj.a, trainer->q_proj.rank * trainer->q_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer->q_proj.b, trainer->q_proj.out_dim * trainer->q_proj.rank); + out_result->v_param_abs_sum = orion_qwen_lora_abs_sum(trainer->v_proj.a, trainer->v_proj.rank * trainer->v_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer->v_proj.b, trainer->v_proj.out_dim * trainer->v_proj.rank); + out_result->predicted_token = -1; + ok = 1; + +fail: + if (ioIn) CFRelease(ioIn); + if (ioNormedIn) CFRelease(ioNormedIn); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + free(hidden_in); free(normed); free(base_q); free(base_v); free(cpu_q); free(cpu_v); free(delta_q); free(delta_v); + free(q_full); free(v_raw); free(attn_cat); free(gated); free(mixer); free(hidden_mid); + free(post_norm); free(mlp_out); free(hidden_out); free(final_norm); free(last_hidden); + free(d_last_hidden); free(d_hidden_out); free(d_post_norm); free(d_hidden_mid); free(d_mixer); + free(d_gated); free(d_attn_cat); free(d_gate_half); free(d_q_full); free(d_v); free(d_normed_from_q); + free(d_normed_from_v); free(d_normed); free(d_hidden_mid_from_post); free(d_post_ln_weight_grad); + free(input_ln); free(post_ln); free(q_proj); free(v_proj); free(o_proj); free(gate_proj); free(up_proj); free(down_proj); + free(final_norm_weight); + orion_qwen35_manifest_free(manifest); + return ok; +} diff --git a/model/configs/gpt2_124m.h b/model/configs/gpt2_124m.h index dbca257..0094c9b 100644 --- a/model/configs/gpt2_124m.h +++ b/model/configs/gpt2_124m.h @@ -15,6 +15,7 @@ static const OrionModelConfig kGPT2_124M = { .hidden_dim = 3072, // 4 * 768 .vocab = 50257, .max_seq = 1024, + .n_kv_head = 12, }; // Inference bucket sizes for ANE prefill diff --git a/model/configs/qwen35_08b.h b/model/configs/qwen35_08b.h new file mode 100644 index 0000000..88e4306 --- /dev/null +++ b/model/configs/qwen35_08b.h @@ -0,0 +1,20 @@ +#ifndef ORION_QWEN35_08B_H +#define ORION_QWEN35_08B_H + +#import "../../core/ane_runtime.h" + +// Qwen3.5-0.8B text-only runtime configuration. +// These values are taken from the Hugging Face config.text_config dump and are +// used only for stage-1 CPU-only port scaffolding. +static const OrionModelConfig kQwen35_08B = { + .n_layer = 24, + .n_head = 8, + .d_model = 1024, + .head_dim = 256, + .hidden_dim = 3584, + .vocab = 248320, + .max_seq = 262144, + .n_kv_head = 2, +}; + +#endif // ORION_QWEN35_08B_H diff --git a/model/configs/qwen35_9b.h b/model/configs/qwen35_9b.h new file mode 100644 index 0000000..2a0c989 --- /dev/null +++ b/model/configs/qwen35_9b.h @@ -0,0 +1,20 @@ +#ifndef ORION_QWEN35_9B_H +#define ORION_QWEN35_9B_H + +#import "../../core/ane_runtime.h" + +// Qwen3.5-9B text-only runtime configuration. +// Values are taken from the local export config dump and used for +// CPU-only / ANE prefill port scaffolding. +static const OrionModelConfig kQwen35_9B = { + .n_layer = 32, + .n_head = 16, + .d_model = 4096, + .head_dim = 256, + .hidden_dim = 12288, + .vocab = 248320, + .max_seq = 262144, + .n_kv_head = 4, +}; + +#endif // ORION_QWEN35_9B_H diff --git a/model/configs/stories110m.h b/model/configs/stories110m.h index 36e8f68..7220866 100644 --- a/model/configs/stories110m.h +++ b/model/configs/stories110m.h @@ -16,6 +16,7 @@ static const OrionModelConfig kStories110M = { .hidden_dim = 2048, // SwiGLU hidden size .vocab = 32000, // Llama2 BPE vocabulary .max_seq = 256, + .n_kv_head = 12, }; // Training constants (from ANEgpt) diff --git a/model/convert/hf_to_blobs_qwen35.py b/model/convert/hf_to_blobs_qwen35.py new file mode 100644 index 0000000..44e06c0 --- /dev/null +++ b/model/convert/hf_to_blobs_qwen35.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +"""Convert Hugging Face Qwen3.5 text weights into Orion BLOBFILE blobs. + +This is a stage-1 converter for Orion Qwen porting: + - understands Qwen3.5 text_config / layer_types + - exports text-only weights and a manifest + - explicitly ignores visual / mtp branches + +It is intentionally conservative. The goal is to unblock loader smoke and +CPU-only inference porting, not to claim full runtime compatibility yet. +""" + +import argparse +import json +import os +from pathlib import Path + +try: + from huggingface_hub import hf_hub_download +except ImportError as exc: # pragma: no cover - import guard + raise SystemExit("pip install huggingface_hub") from exc + +try: + from transformers import AutoConfig, AutoModelForCausalLM +except ImportError as exc: # pragma: no cover - import guard + raise SystemExit("pip install transformers torch") from exc + +from hf_to_blobs_gpt2 import convert_tensor_to_blob, convert_tensor_transposed + + +TEXT_PREFIX = "model" +VISUAL_PREFIX = "visual" +MTP_PREFIX = "mtp" + + +def normalize_index_name(name): + if name.startswith("model.language_model."): + return "model." + name[len("model.language_model."):] + if name.startswith("language_model.model."): + return "model." + name[len("language_model.model."):] + if name.startswith("language_model.lm_head."): + return "lm_head." + name[len("language_model.lm_head."):] + return name + + +def is_local_model_path(model_name): + return Path(model_name).expanduser().exists() + + +def resolve_local_model_path(model_name): + return Path(model_name).expanduser().resolve() + + +def output_entry(layer_idx, blob_name, hf_name, shape, transpose, alias_of=None): + if layer_idx is None: + rel_path = f"model/{blob_name}.bin" + else: + rel_path = f"layer{layer_idx}/{blob_name}.bin" + return { + "path": rel_path, + "hf_name": hf_name, + "shape": list(shape) if shape is not None else None, + "transpose": bool(transpose), + "alias_of": alias_of, + } + + +def runtime_from_config(cfg): + tc = cfg.text_config if hasattr(cfg, "text_config") else cfg + rope_parameters = getattr(tc, "rope_parameters", {}) or {} + return { + "source": "text_config" if tc is not cfg else "root", + "hidden_size": getattr(tc, "hidden_size", None), + "num_hidden_layers": getattr(tc, "num_hidden_layers", None), + "num_attention_heads": getattr(tc, "num_attention_heads", None), + "num_key_value_heads": getattr(tc, "num_key_value_heads", None), + "intermediate_size": getattr(tc, "intermediate_size", None), + "vocab_size": getattr(tc, "vocab_size", None), + "max_position_embeddings": getattr(tc, "max_position_embeddings", None), + "rms_norm_eps": getattr(tc, "rms_norm_eps", None), + "head_dim": getattr(tc, "head_dim", None), + "tie_word_embeddings": getattr(tc, "tie_word_embeddings", None), + "rope_parameters": rope_parameters, + "layer_types": list(getattr(tc, "layer_types", [])), + } + + +def flatten_text_config(cfg): + tc = cfg.text_config if hasattr(cfg, "text_config") else None + if tc is None: + return cfg + + for key, value in tc.to_dict().items(): + if key.startswith("_"): + continue + setattr(cfg, key, value) + return cfg + + +def build_expected_entries(cfg): + runtime = runtime_from_config(cfg) + layer_types = runtime["layer_types"] + entries = [ + output_entry(None, "embed_tokens", f"{TEXT_PREFIX}.embed_tokens.weight", (runtime["vocab_size"], runtime["hidden_size"]), False), + output_entry(None, "final_norm", f"{TEXT_PREFIX}.norm.weight", (runtime["hidden_size"],), False), + output_entry(None, "lm_head", "lm_head.weight", (runtime["vocab_size"], runtime["hidden_size"]), False), + ] + + for idx, layer_type in enumerate(layer_types): + base = f"{TEXT_PREFIX}.layers.{idx}" + entries.append(output_entry(idx, "input_layernorm", f"{base}.input_layernorm.weight", (runtime["hidden_size"],), False)) + entries.append(output_entry(idx, "post_attention_layernorm", f"{base}.post_attention_layernorm.weight", (runtime["hidden_size"],), False)) + entries.append(output_entry(idx, "mlp_gate_proj", f"{base}.mlp.gate_proj.weight", (runtime["intermediate_size"], runtime["hidden_size"]), True)) + entries.append(output_entry(idx, "mlp_up_proj", f"{base}.mlp.up_proj.weight", (runtime["intermediate_size"], runtime["hidden_size"]), True)) + entries.append(output_entry(idx, "mlp_down_proj", f"{base}.mlp.down_proj.weight", (runtime["hidden_size"], runtime["intermediate_size"]), True)) + + if layer_type == "full_attention": + entries.extend( + [ + output_entry(idx, "self_attn_q_proj", f"{base}.self_attn.q_proj.weight", (runtime["hidden_size"], runtime["hidden_size"]), True), + output_entry(idx, "self_attn_k_proj", f"{base}.self_attn.k_proj.weight", (runtime["num_key_value_heads"] * runtime["head_dim"], runtime["hidden_size"]), True), + output_entry(idx, "self_attn_v_proj", f"{base}.self_attn.v_proj.weight", (runtime["num_key_value_heads"] * runtime["head_dim"], runtime["hidden_size"]), True), + output_entry(idx, "self_attn_o_proj", f"{base}.self_attn.o_proj.weight", (runtime["hidden_size"], runtime["hidden_size"]), True), + output_entry(idx, "self_attn_q_norm", f"{base}.self_attn.q_norm.weight", (runtime["head_dim"],), False), + output_entry(idx, "self_attn_k_norm", f"{base}.self_attn.k_norm.weight", (runtime["head_dim"],), False), + ] + ) + elif layer_type == "linear_attention": + entries.extend( + [ + output_entry(idx, "linear_attn_in_proj_qkv", f"{base}.linear_attn.in_proj_qkv.weight", None, True), + output_entry(idx, "linear_attn_in_proj_z", f"{base}.linear_attn.in_proj_z.weight", None, True), + output_entry(idx, "linear_attn_in_proj_a", f"{base}.linear_attn.in_proj_a.weight", None, True), + output_entry(idx, "linear_attn_in_proj_b", f"{base}.linear_attn.in_proj_b.weight", None, True), + output_entry(idx, "linear_attn_out_proj", f"{base}.linear_attn.out_proj.weight", None, True), + output_entry(idx, "linear_attn_norm", f"{base}.linear_attn.norm.weight", None, False), + output_entry(idx, "linear_attn_dt_bias", f"{base}.linear_attn.dt_bias", None, False), + output_entry(idx, "linear_attn_a_log", f"{base}.linear_attn.A_log", None, False), + output_entry(idx, "linear_attn_conv1d", f"{base}.linear_attn.conv1d.weight", None, False), + ] + ) + else: + raise ValueError(f"Unsupported layer type: {layer_type}") + + return runtime, entries + + +def load_index(model_name): + if is_local_model_path(model_name): + index_path = resolve_local_model_path(model_name) / "model.safetensors.index.json" + if not index_path.exists(): + raise FileNotFoundError(f"Local model index not found: {index_path}") + else: + index_path = hf_hub_download(model_name, "model.safetensors.index.json") + with open(index_path, "r", encoding="utf-8") as f: + return json.load(f) + + +def build_manifest(model_name, output_dir): + cfg = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + runtime, entries = build_expected_entries(cfg) + index_data = load_index(model_name) + raw_weight_map = index_data["weight_map"] + weight_map = {normalize_index_name(k): v for k, v in raw_weight_map.items()} + available_keys = set(weight_map.keys()) + metadata = index_data.get("metadata", {}) or {} + total_size = metadata.get("total_size") + missing = [] + present = [] + for entry in entries: + if entry["hf_name"] in available_keys: + present.append(entry) + else: + missing.append(entry["hf_name"]) + + # Qwen3.5 ties lm_head to embed_tokens, so lm_head.weight may be absent + # from the safetensors index even though the runtime should treat it as + # the same tensor. + if runtime.get("tie_word_embeddings") and "lm_head.weight" in missing: + missing = [name for name in missing if name != "lm_head.weight"] + present.append( + output_entry( + None, + "lm_head", + "lm_head.weight", + (runtime["vocab_size"], runtime["hidden_size"]), + False, + alias_of=f"{TEXT_PREFIX}.embed_tokens.weight", + ) + ) + + manifest = { + "status": "PASS_QWEN35_MANIFEST_ONLY" if not missing else "BLOCKED_QWEN35_MANIFEST_MISSING", + "model": model_name, + "runtime": runtime, + "present_entries": present, + "missing_hf_names": missing, + "weight_map_shards": weight_map, + "hf_total_size_bytes": total_size, + "ignored_prefixes": [VISUAL_PREFIX, MTP_PREFIX], + "notes": [ + "Stage-1 converter is text-only.", + "Visual and mtp branches are intentionally excluded.", + "ANE runtime support is not implied by this manifest.", + ], + } + + os.makedirs(output_dir, exist_ok=True) + manifest_path = os.path.join(output_dir, "manifest.json") + with open(manifest_path, "w", encoding="utf-8") as f: + json.dump(manifest, f, ensure_ascii=False, indent=2) + return manifest_path, manifest + + +def apply_actual_shapes_from_state(manifest, state): + updated_entries = [] + for entry in manifest["present_entries"]: + tensor_name = entry["alias_of"] or entry["hf_name"] + tensor = state[tensor_name] + updated = dict(entry) + updated["shape"] = list(tensor.shape) + updated["resolved_hf_name"] = tensor_name + updated_entries.append(updated) + manifest["present_entries"] = updated_entries + return manifest + + +def export_qwen35(model_name, output_dir): + manifest_path, manifest = build_manifest(model_name, output_dir) + if manifest["missing_hf_names"]: + raise RuntimeError(f"Missing expected tensors: {manifest['missing_hf_names'][:5]}") + + print(f"Loading model {model_name} for text-only export...") + cfg = flatten_text_config(AutoConfig.from_pretrained(model_name, trust_remote_code=True)) + model = AutoModelForCausalLM.from_pretrained(model_name, config=cfg, trust_remote_code=True) + state = model.state_dict() + manifest = apply_actual_shapes_from_state(manifest, state) + with open(manifest_path, "w", encoding="utf-8") as f: + json.dump(manifest, f, ensure_ascii=False, indent=2) + + total_bytes = 0 + file_count = 0 + saved = [] + + for entry in manifest["present_entries"]: + tensor_name = entry["alias_of"] or entry["hf_name"] + tensor = state[tensor_name].detach().cpu().float().numpy() + out_path = os.path.join(output_dir, entry["path"]) + Path(out_path).parent.mkdir(parents=True, exist_ok=True) + if entry["transpose"] and tensor.ndim == 2: + size = convert_tensor_transposed(tensor, out_path) + else: + size = convert_tensor_to_blob(tensor, out_path) + saved.append({**entry, "resolved_hf_name": tensor_name, "bytes_fp16": size}) + total_bytes += size + file_count += 1 + print(f" {entry['path']}: {tuple(tensor.shape)} -> {size} bytes fp16") + + export_summary = { + "status": "PASS_QWEN35_TEXT_ONLY_EXPORT", + "model": model_name, + "manifest_path": manifest_path, + "file_count": file_count, + "total_bytes_fp16": total_bytes, + "saved_entries": saved, + } + summary_path = os.path.join(output_dir, "export_summary.json") + with open(summary_path, "w", encoding="utf-8") as f: + json.dump(export_summary, f, ensure_ascii=False, indent=2) + print(json.dumps(export_summary, ensure_ascii=False, indent=2)) + return summary_path + + +def main(): + parser = argparse.ArgumentParser(description="Convert Qwen3.5 text weights to Orion BLOBFILE format") + parser.add_argument("--model", default="Qwen/Qwen3.5-0.8B", help="Hugging Face model name or local export directory") + parser.add_argument("--output", required=True, help="Output directory") + parser.add_argument("--manifest-only", action="store_true", help="Only write manifest.json without exporting weights") + args = parser.parse_args() + + if args.manifest_only: + manifest_path, manifest = build_manifest(args.model, args.output) + print(f"MANIFEST_PATH={manifest_path}") + print(json.dumps(manifest, ensure_ascii=False, indent=2)) + else: + export_qwen35(args.model, args.output) + + +if __name__ == "__main__": + main() diff --git a/model/weight_loader.h b/model/weight_loader.h index 6f08323..205e5e4 100644 --- a/model/weight_loader.h +++ b/model/weight_loader.h @@ -2,6 +2,7 @@ #define ORION_WEIGHT_LOADER_H #import +#include /// GPT-2 124M layer weights (fp32, loaded from BLOBFILE blobs). typedef struct { @@ -37,6 +38,32 @@ typedef struct { float *ln_f_b; // [d_model] final LayerNorm beta } OrionGPT2Weights; +/// Lightweight Qwen3.5 text-only runtime metadata loaded from manifest.json. +/// This is a stage-1 porting structure used to validate that converted blobs +/// and runtime dimensions are coherent before full CPU-only inference exists. +typedef struct { + int n_layer; + int d_model; + int d_ff; + int n_head; + int n_kv_head; + int head_dim; + int vocab; + int max_seq; + int tie_word_embeddings; + float rope_theta; + float partial_rotary_factor; + int rotary_dim; + int n_entries; + int n_alias_entries; + int n_linear_layers; + int n_full_layers; + char *blob_dir; + char *manifest_path; +} OrionQwen35Manifest; + +typedef struct OrionBlobRowReader OrionBlobRowReader; + /// Load GPT-2 weights from BLOBFILE directory. /// @param blob_dir Path to directory containing blob files (e.g., model/blobs/gpt2_124m/) /// @return Loaded weights, or NULL on failure. Caller must free with orion_gpt2_weights_free. @@ -45,4 +72,38 @@ OrionGPT2Weights* orion_gpt2_weights_load(const char* blob_dir); /// Free loaded weights. void orion_gpt2_weights_free(OrionGPT2Weights* w); +/// Load stage-1 Qwen3.5 runtime metadata from manifest.json. +/// @param blob_dir Path to converted blob directory that contains manifest.json. +/// @return Loaded manifest metadata, or NULL on failure. Caller must free with +/// orion_qwen35_manifest_free. +OrionQwen35Manifest* orion_qwen35_manifest_load(const char* blob_dir); + +/// Free loaded Qwen3.5 manifest metadata. +void orion_qwen35_manifest_free(OrionQwen35Manifest* manifest); + +/// Return the number of fp16 elements stored in a BLOBFILE. +/// Returns -1 on error. +int orion_blob_element_count(const char* path); + +/// Read a contiguous fp16 row from a BLOBFILE and convert it to fp32. +/// The tensor is treated as a row-major 2D matrix with width `row_width`. +/// Vectors can be read by passing row_index=0 and row_width=. +/// Returns 1 on success, 0 on failure. +int orion_read_blob_row_f32(const char* path, int row_index, int row_width, float* out_row); + +/// Open a buffered row reader for repeated fp16 -> fp32 row access. +/// Returns NULL on failure. +OrionBlobRowReader* orion_blob_row_reader_open(const char* path, int row_width); + +/// Read a row via an open row reader. +/// Returns 1 on success, 0 on failure. +int orion_blob_row_reader_read_f32(OrionBlobRowReader* reader, int row_index, float* out_row); + +/// Close a row reader and release its buffers. +void orion_blob_row_reader_close(OrionBlobRowReader* reader); + +/// Read a full BLOBFILE tensor into fp32 with an exact expected element count. +/// Returns heap-allocated fp32 data on success, or NULL on failure. +float* orion_read_blob_f32_exact(const char* path, int count); + #endif // ORION_WEIGHT_LOADER_H diff --git a/model/weight_loader.m b/model/weight_loader.m index afb891a..1ac81bb 100644 --- a/model/weight_loader.m +++ b/model/weight_loader.m @@ -9,6 +9,16 @@ #pragma mark - BLOBFILE Reader +struct OrionBlobRowReader { + FILE *file; + int row_width; + int row_count; + int cache_row_start; + int cache_row_count; + int cache_capacity_rows; + _Float16 *fp16_cache; +}; + /// Read a BLOBFILE and return fp32 array. Caller must free. /// @param path Path to .bin blob file /// @param count Expected number of elements (for validation) @@ -50,6 +60,307 @@ return fp32; } +float* orion_read_blob_f32_exact(const char* path, int count) { + return read_blob_f32(path, count); +} + +int orion_blob_element_count(const char* path) { + FILE *f = fopen(path, "rb"); + if (!f) { + fprintf(stderr, "weight_loader: cannot open %s\n", path); + return -1; + } + fseek(f, 0, SEEK_END); + long file_size = ftell(f); + fclose(f); + if (file_size < 128) { + fprintf(stderr, "weight_loader: %s: invalid file size %ld\n", path, file_size); + return -1; + } + return (int)((file_size - 128) / 2); +} + +int orion_read_blob_row_f32(const char* path, int row_index, int row_width, float* out_row) { + if (!path || row_index < 0 || row_width <= 0 || !out_row) { + fprintf(stderr, "weight_loader: invalid row read request\n"); + return 0; + } + + int total_count = orion_blob_element_count(path); + if (total_count < 0) { + return 0; + } + if ((total_count % row_width) != 0) { + fprintf(stderr, "weight_loader: %s: row_width %d does not divide %d elements\n", + path, row_width, total_count); + return 0; + } + + int row_count = total_count / row_width; + if (row_index >= row_count) { + fprintf(stderr, "weight_loader: %s: row_index %d out of range (rows=%d)\n", + path, row_index, row_count); + return 0; + } + + FILE *f = fopen(path, "rb"); + if (!f) { + fprintf(stderr, "weight_loader: cannot open %s\n", path); + return 0; + } + + long offset = 128L + ((long)row_index * (long)row_width * 2L); + if (fseek(f, offset, SEEK_SET) != 0) { + fprintf(stderr, "weight_loader: failed to seek %s\n", path); + fclose(f); + return 0; + } + + _Float16 *fp16 = (_Float16 *)malloc((size_t)row_width * sizeof(_Float16)); + if (!fp16) { + fclose(f); + return 0; + } + size_t nread = fread(fp16, sizeof(_Float16), (size_t)row_width, f); + fclose(f); + if (nread != (size_t)row_width) { + fprintf(stderr, "weight_loader: short read on %s\n", path); + free(fp16); + return 0; + } + + for (int i = 0; i < row_width; i++) { + out_row[i] = (float)fp16[i]; + } + free(fp16); + return 1; +} + +OrionBlobRowReader* orion_blob_row_reader_open(const char* path, int row_width) { + if (!path || row_width <= 0) { + fprintf(stderr, "weight_loader: invalid row reader open request\n"); + return NULL; + } + + int total_count = orion_blob_element_count(path); + if (total_count < 0) { + return NULL; + } + if ((total_count % row_width) != 0) { + fprintf(stderr, "weight_loader: %s: row_width %d does not divide %d elements\n", + path, row_width, total_count); + return NULL; + } + + FILE *f = fopen(path, "rb"); + if (!f) { + fprintf(stderr, "weight_loader: cannot open %s\n", path); + return NULL; + } + setvbuf(f, NULL, _IOFBF, 1 << 20); + + OrionBlobRowReader *reader = calloc(1, sizeof(OrionBlobRowReader)); + if (!reader) { + fclose(f); + return NULL; + } + + const int row_bytes = row_width * (int)sizeof(_Float16); + const int target_cache_bytes = 1 << 20; + int cache_capacity_rows = target_cache_bytes / row_bytes; + if (cache_capacity_rows < 1) cache_capacity_rows = 1; + if (cache_capacity_rows > 256) cache_capacity_rows = 256; + if (cache_capacity_rows > (total_count / row_width)) { + cache_capacity_rows = total_count / row_width; + } + + reader->fp16_cache = (_Float16 *)malloc((size_t)row_width * (size_t)cache_capacity_rows * sizeof(_Float16)); + if (!reader->fp16_cache) { + fclose(f); + free(reader); + return NULL; + } + + reader->file = f; + reader->row_width = row_width; + reader->row_count = total_count / row_width; + reader->cache_row_start = -1; + reader->cache_row_count = 0; + reader->cache_capacity_rows = cache_capacity_rows; + return reader; +} + +int orion_blob_row_reader_read_f32(OrionBlobRowReader* reader, int row_index, float* out_row) { + if (!reader || !reader->file || !out_row || row_index < 0 || row_index >= reader->row_count) { + fprintf(stderr, "weight_loader: invalid reader row request idx=%d\n", row_index); + return 0; + } + + const int cache_row_end = reader->cache_row_start + reader->cache_row_count; + if (!(reader->cache_row_start >= 0 && + row_index >= reader->cache_row_start && + row_index < cache_row_end)) { + const int block_row_start = (row_index / reader->cache_capacity_rows) * reader->cache_capacity_rows; + const int block_row_count = (reader->row_count - block_row_start) < reader->cache_capacity_rows + ? (reader->row_count - block_row_start) + : reader->cache_capacity_rows; + + long offset = 128L + ((long)block_row_start * (long)reader->row_width * 2L); + if (fseek(reader->file, offset, SEEK_SET) != 0) { + fprintf(stderr, "weight_loader: row reader seek failed\n"); + return 0; + } + + const size_t nread = fread(reader->fp16_cache, + sizeof(_Float16), + (size_t)reader->row_width * (size_t)block_row_count, + reader->file); + if (nread != (size_t)reader->row_width * (size_t)block_row_count) { + fprintf(stderr, "weight_loader: row reader short read\n"); + return 0; + } + + reader->cache_row_start = block_row_start; + reader->cache_row_count = block_row_count; + } + + const int local_row = row_index - reader->cache_row_start; + const _Float16 *fp16_row = reader->fp16_cache + ((size_t)local_row * (size_t)reader->row_width); + for (int i = 0; i < reader->row_width; i++) { + out_row[i] = (float)fp16_row[i]; + } + return 1; +} + +void orion_blob_row_reader_close(OrionBlobRowReader* reader) { + if (!reader) return; + if (reader->file) fclose(reader->file); + free(reader->fp16_cache); + free(reader); +} + +#pragma mark - Qwen3.5 Manifest Loader + +static int json_int(NSDictionary *dict, NSString *key) { + id value = dict[key]; + if (![value respondsToSelector:@selector(intValue)]) { + return 0; + } + return [value intValue]; +} + +static float json_float(NSDictionary *dict, NSString *key, float fallback) { + id value = dict[key]; + if (![value respondsToSelector:@selector(floatValue)]) { + return fallback; + } + return [value floatValue]; +} + +OrionQwen35Manifest* orion_qwen35_manifest_load(const char* blob_dir) { + if (!blob_dir) { + fprintf(stderr, "weight_loader: qwen35 manifest load got NULL blob_dir\n"); + return NULL; + } + + @autoreleasepool { + NSString *blobDir = [NSString stringWithUTF8String:blob_dir]; + NSString *manifestPath = [blobDir stringByAppendingPathComponent:@"manifest.json"]; + NSData *data = [NSData dataWithContentsOfFile:manifestPath]; + if (!data) { + fprintf(stderr, "weight_loader: cannot read %s\n", manifestPath.UTF8String); + return NULL; + } + + NSError *error = nil; + id root = [NSJSONSerialization JSONObjectWithData:data options:0 error:&error]; + if (![root isKindOfClass:[NSDictionary class]]) { + fprintf(stderr, "weight_loader: invalid manifest json: %s\n", + error.localizedDescription.UTF8String); + return NULL; + } + + NSDictionary *manifest = (NSDictionary *)root; + NSString *status = manifest[@"status"]; + if (![status isEqualToString:@"PASS_QWEN35_MANIFEST_ONLY"] && + ![status isEqualToString:@"PASS_QWEN35_TEXT_ONLY_EXPORT"]) { + fprintf(stderr, "weight_loader: unexpected manifest status %s\n", + status.UTF8String); + return NULL; + } + + NSArray *missing = manifest[@"missing_hf_names"]; + if ([missing isKindOfClass:[NSArray class]] && [missing count] > 0) { + fprintf(stderr, "weight_loader: manifest still has missing tensors (%lu)\n", + (unsigned long)[missing count]); + return NULL; + } + + NSDictionary *runtime = manifest[@"runtime"]; + NSArray *entries = manifest[@"present_entries"]; + if (![runtime isKindOfClass:[NSDictionary class]] || + ![entries isKindOfClass:[NSArray class]]) { + fprintf(stderr, "weight_loader: manifest missing runtime/present_entries\n"); + return NULL; + } + + OrionQwen35Manifest *out = calloc(1, sizeof(OrionQwen35Manifest)); + out->n_layer = json_int(runtime, @"num_hidden_layers"); + out->d_model = json_int(runtime, @"hidden_size"); + out->d_ff = json_int(runtime, @"intermediate_size"); + out->n_head = json_int(runtime, @"num_attention_heads"); + out->n_kv_head = json_int(runtime, @"num_key_value_heads"); + out->head_dim = json_int(runtime, @"head_dim"); + out->vocab = json_int(runtime, @"vocab_size"); + out->max_seq = json_int(runtime, @"max_position_embeddings"); + out->tie_word_embeddings = json_int(runtime, @"tie_word_embeddings"); + NSDictionary *ropeParams = runtime[@"rope_parameters"]; + out->rope_theta = 10000000.0f; + out->partial_rotary_factor = 0.25f; + if ([ropeParams isKindOfClass:[NSDictionary class]]) { + out->rope_theta = json_float(ropeParams, @"rope_theta", out->rope_theta); + out->partial_rotary_factor = json_float(ropeParams, @"partial_rotary_factor", out->partial_rotary_factor); + } + out->rotary_dim = (int)(out->head_dim * out->partial_rotary_factor); + if (out->rotary_dim % 2 != 0) { + out->rotary_dim -= 1; + } + out->n_entries = (int)[entries count]; + out->blob_dir = strdup(blob_dir); + out->manifest_path = strdup(manifestPath.UTF8String); + + NSArray *layerTypes = runtime[@"layer_types"]; + if ([layerTypes isKindOfClass:[NSArray class]]) { + for (id value in layerTypes) { + if (![value isKindOfClass:[NSString class]]) continue; + NSString *layerType = (NSString *)value; + if ([layerType isEqualToString:@"linear_attention"]) { + out->n_linear_layers += 1; + } else if ([layerType isEqualToString:@"full_attention"]) { + out->n_full_layers += 1; + } + } + } + + for (id value in entries) { + if (![value isKindOfClass:[NSDictionary class]]) continue; + NSDictionary *entry = (NSDictionary *)value; + id aliasOf = entry[@"alias_of"]; + if ([aliasOf isKindOfClass:[NSString class]] && [(NSString *)aliasOf length] > 0) { + out->n_alias_entries += 1; + } + } + + if (out->n_layer <= 0 || out->d_model <= 0 || out->vocab <= 0) { + fprintf(stderr, "weight_loader: invalid qwen35 runtime dims from manifest\n"); + orion_qwen35_manifest_free(out); + return NULL; + } + + return out; + } +} + #pragma mark - Load OrionGPT2Weights* orion_gpt2_weights_load(const char* blob_dir) { @@ -59,6 +370,23 @@ int d_ff = 3072; // 4 * d_model int vocab = 50257; int max_seq = 1024; + char path[1024]; + + snprintf(path, sizeof(path), "%s/wte.bin", blob_dir); + int wte_count = orion_blob_element_count(path); + if (wte_count <= 0 || (wte_count % d_model) != 0) { + fprintf(stderr, "weight_loader: failed to infer vocab from %s\n", path); + return NULL; + } + vocab = wte_count / d_model; + + snprintf(path, sizeof(path), "%s/wpe.bin", blob_dir); + int wpe_count = orion_blob_element_count(path); + if (wpe_count <= 0 || (wpe_count % d_model) != 0) { + fprintf(stderr, "weight_loader: failed to infer max_seq from %s\n", path); + return NULL; + } + max_seq = wpe_count / d_model; OrionGPT2Weights *w = (OrionGPT2Weights *)calloc(1, sizeof(OrionGPT2Weights)); w->n_layer = n_layer; @@ -66,9 +394,6 @@ w->d_ff = d_ff; w->vocab = vocab; w->max_seq = max_seq; - - char path[1024]; - // Token embedding [vocab, d_model] snprintf(path, sizeof(path), "%s/wte.bin", blob_dir); w->wte = read_blob_f32(path, vocab * d_model); @@ -150,3 +475,10 @@ void orion_gpt2_weights_free(OrionGPT2Weights* w) { } free(w); } + +void orion_qwen35_manifest_free(OrionQwen35Manifest* manifest) { + if (!manifest) return; + free(manifest->blob_dir); + free(manifest->manifest_path); + free(manifest); +} diff --git a/scripts/orion_q_share_manifest.txt b/scripts/orion_q_share_manifest.txt new file mode 100644 index 0000000..082c594 --- /dev/null +++ b/scripts/orion_q_share_manifest.txt @@ -0,0 +1,48 @@ +# Orion-Q share manifest +# Paths are relative to tmp/Orion and may use glob patterns. + +.gitignore +LICENSE +Makefile +README.md +README_ORION_Q.md +requirements.txt + +docs/orion_q/*.md + +apps/cli/commands/infer.m + +compiler/builder.c +compiler/model_config.h +compiler/frontends/qwen35_prefill.c +compiler/frontends/qwen35_prefill.h + +core/ane_runtime.h +core/ane_runtime.m +core/mil_builder.m +core/model_registry.m + +kernels/inference/qwen_cpu_ops.h +kernels/inference/qwen_cpu_ops.m +kernels/training/qwen_lora_cpu_ops.h +kernels/training/qwen_lora_cpu_ops.m +kernels/training/qwen_lora_train.h +kernels/training/qwen_lora_train.m + +model/configs/gpt2_124m.h +model/configs/stories110m.h +model/configs/qwen35_08b.h +model/configs/qwen35_9b.h +model/convert/hf_to_blobs_qwen35.py +model/weight_loader.h +model/weight_loader.m + +scripts/orion_q_share_manifest.txt +scripts/prepare_orion_q_share.py + +tokenizer/gpt2_bpe.h +tokenizer/gpt2_bpe.m + +tests/test_compiler_equiv.m +tests/test_qwen35*.m +tests/test_qwen35_9b*.m diff --git a/scripts/prepare_orion_q_share.py b/scripts/prepare_orion_q_share.py new file mode 100644 index 0000000..0185c09 --- /dev/null +++ b/scripts/prepare_orion_q_share.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import shutil +from pathlib import Path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Prepare a clean Orion-Q share bundle.") + parser.add_argument( + "--output", + type=Path, + default=None, + help="Output directory. Default: /build/orion_q_share", + ) + parser.add_argument( + "--manifest", + type=Path, + default=None, + help="Manifest file. Default: tmp/Orion/scripts/orion_q_share_manifest.txt", + ) + parser.add_argument( + "--clean", + action="store_true", + help="Remove the existing output directory before copying.", + ) + return parser.parse_args() + + +def read_manifest(manifest_path: Path) -> list[str]: + patterns: list[str] = [] + for raw_line in manifest_path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + patterns.append(line) + return patterns + + +def expand_patterns(root: Path, patterns: list[str]) -> tuple[list[Path], list[str]]: + matched: set[Path] = set() + missing: list[str] = [] + for pattern in patterns: + hits = [path for path in root.glob(pattern) if path.is_file()] + if not hits: + missing.append(pattern) + continue + matched.update(hits) + return sorted(matched), missing + + +def copy_file(src: Path, dst: Path) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dst) + + +def write_manifest(bundle_root: Path, copied_files: list[str], missing_patterns: list[str]) -> None: + lines = [ + "# Orion-Q share bundle", + "", + "## Copied files", + *copied_files, + ] + if missing_patterns: + lines.extend([ + "", + "## Manifest patterns with no matches", + *missing_patterns, + ]) + (bundle_root / "MANIFEST.generated.txt").write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def main() -> int: + args = parse_args() + script_path = Path(__file__).resolve() + orion_root = script_path.parents[1] + manifest_path = args.manifest or (orion_root / "scripts" / "orion_q_share_manifest.txt") + output_root = args.output or (orion_root / "build" / "orion_q_share") + bundle_root = output_root / "Orion-Q" + + if args.clean and output_root.exists(): + shutil.rmtree(output_root) + output_root.mkdir(parents=True, exist_ok=True) + + patterns = read_manifest(manifest_path) + source_files, missing_patterns = expand_patterns(orion_root, patterns) + + copied: list[str] = [] + for src in source_files: + rel = src.relative_to(orion_root) + copy_file(src, bundle_root / rel) + copied.append(str(rel)) + + write_manifest(bundle_root, copied, missing_patterns) + + print(f"Prepared Orion-Q share bundle: {bundle_root}") + print(f"Copied files: {len(copied)}") + if missing_patterns: + print(f"Unmatched manifest entries: {len(missing_patterns)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_compiler_equiv.m b/tests/test_compiler_equiv.m index dc817a6..a032361 100644 --- a/tests/test_compiler_equiv.m +++ b/tests/test_compiler_equiv.m @@ -142,7 +142,7 @@ static bool test_stories_fwd_ffn(void) { NSString* mil = orion_codegen_mil(g, "main"); orion_graph_free(g); - return mil != nil && [mil containsString:@"sigmoid("]; // SiLU decomposition + return mil != nil && [mil containsString:@"tanh("] && [mil containsString:@"_sig"]; } // Test: Stories ffnBwd generates valid MIL diff --git a/tests/test_qwen35_24layer_logits_cpu_smoke.m b/tests/test_qwen35_24layer_logits_cpu_smoke.m new file mode 100644 index 0000000..569cb4a --- /dev/null +++ b/tests/test_qwen35_24layer_logits_cpu_smoke.m @@ -0,0 +1,244 @@ +#import +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static int argmax_vocab_from_tied_embedding(const char *blob_dir, + const float *hidden, + int d_model, + int vocab, + int *top_id, + float *top_logit) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + int best_id = -1; + float best_logit = -INFINITY; + for (int tok = 0; tok < vocab; tok++) { + if (!orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) { + dot += hidden[i] * row[i]; + } + if (dot > best_logit) { + best_logit = dot; + best_id = tok; + } + } + + free(row); + *top_id = best_id; + *top_logit = best_logit; + return 1; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int seq_len = 2; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int num_k_heads = 16; + const int num_v_heads = 16; + const int head_k_dim = 128; + const int head_v_dim = 128; + const int value_dim = num_v_heads * head_v_dim; + const int conv_kernel = 4; + + float *hidden = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = NULL; + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch || !last_hidden) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(embed_path, 0, d_model, hidden + 0 * d_model) || + !orion_read_blob_row_f32(embed_path, 1, d_model, hidden + 1 * d_model)) { + fprintf(stderr, "FAIL: failed to read token embeddings\n"); + goto fail; + } + + for (int layer_idx = 0; layer_idx < manifest->n_layer; layer_idx++) { + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + fprintf(stderr, "FAIL: missing layer %d norm/mlp tensors\n", layer_idx); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + + memset(mixer_out, 0, (size_t)seq_len * d_model * sizeof(float)); + + int is_full_layer = ((layer_idx + 1) % 4 == 0); + if (is_full_layer) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + fprintf(stderr, "FAIL: missing full-attention tensors for layer %d\n", layer_idx); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", (value_dim + value_dim + value_dim) * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", (value_dim + value_dim + value_dim) * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + fprintf(stderr, "FAIL: missing linear-attention tensors for layer %d\n", layer_idx); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) { + hidden[i] += mixer_out[i]; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + + for (int i = 0; i < seq_len * d_model; i++) { + hidden[i] += mlp_out[i]; + } + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) { + fprintf(stderr, "FAIL: missing final norm\n"); + goto fail; + } + + orion_qwen_cpu_rmsnorm(hidden + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, last_hidden); + + int top_id = -1; + float top_logit = -INFINITY; + if (!argmax_vocab_from_tied_embedding(blob_dir, last_hidden, d_model, manifest->vocab, &top_id, &top_logit)) { + fprintf(stderr, "FAIL: failed vocab top1 scan\n"); + goto fail; + } + + double hidden_abs = abs_sum(hidden, seq_len * d_model); + double last_hidden_abs = abs_sum(last_hidden, d_model); + if (hidden_abs <= 0.0 || last_hidden_abs <= 0.0 || top_id < 0 || !isfinite(top_logit)) { + fprintf(stderr, "FAIL: expected non-zero logits path\n"); + goto fail; + } + + printf("PASS: qwen35 24layer logits cpu smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" seq_len=%d\n", seq_len); + printf(" layers_run=%d\n", manifest->n_layer); + printf(" hidden_abs_sum=%.6f\n", hidden_abs); + printf(" last_hidden_abs_sum=%.6f\n", last_hidden_abs); + printf(" top_token_id=%d\n", top_id); + printf(" top_token_logit=%.6f\n", top_logit); + printf(" next_blocker=%s\n", "decode loop and tokenizer parity validation"); + + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + free(last_hidden); + free(final_norm); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + free(last_hidden); + free(final_norm); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_9b_bridge_stage_diff.m b/tests/test_qwen35_9b_bridge_stage_diff.m new file mode 100644 index 0000000..63f59a3 --- /dev/null +++ b/tests/test_qwen35_9b_bridge_stage_diff.m @@ -0,0 +1,3047 @@ +#import +#import +#import +#include +#include +#include +#include +#include +#include +#include "compiler/builder.h" +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#import "../model/weight_loader.h" +#import "../tokenizer/gpt2_bpe.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static NSDictionary* load_json(NSString* path) { + NSData* data = [NSData dataWithContentsOfFile:path]; + if (!data) return nil; + return [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; +} + +static int file_exists(const char *path) { + return access(path, F_OK) == 0; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static double mean_abs_diff(const float *a, const float *b, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)a[i] - (double)b[i]); + return total / (double)n; +} + +static double max_abs_diff(const float *a, const float *b, int n) { + double best = 0.0; + for (int i = 0; i < n; i++) { + double d = fabs((double)a[i] - (double)b[i]); + if (d > best) best = d; + } + return best; +} + +typedef struct { + long long lhs_size; + long long rhs_size; + long long first_diff_offset; + int equal; +} OrionByteCompareResult; + +static OrionByteCompareResult compare_nsdata(NSData *lhs, NSData *rhs) { + OrionByteCompareResult result = { + .lhs_size = lhs ? (long long)lhs.length : -1, + .rhs_size = rhs ? (long long)rhs.length : -1, + .first_diff_offset = -1, + .equal = 0, + }; + if (!lhs || !rhs) return result; + if ([lhs isEqualToData:rhs]) { + result.equal = 1; + return result; + } + + const unsigned char *lhs_bytes = (const unsigned char *)lhs.bytes; + const unsigned char *rhs_bytes = (const unsigned char *)rhs.bytes; + NSUInteger min_len = lhs.length < rhs.length ? lhs.length : rhs.length; + for (NSUInteger i = 0; i < min_len; i++) { + if (lhs_bytes[i] != rhs_bytes[i]) { + result.first_diff_offset = (long long)i; + return result; + } + } + result.first_diff_offset = (long long)min_len; + return result; +} + +static NSData *load_blob_payload_data(NSString *path) { + NSData *blob = [NSData dataWithContentsOfFile:path]; + if (!blob || blob.length < 128) return nil; + return [blob subdataWithRange:NSMakeRange(128, blob.length - 128)]; +} + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void read_ane_surface_prefix_token_major(IOSurfaceRef s, int channels, int seq_len, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[t * channels + c]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static double mean_abs_diff_q_half(const float *a, const float *b, int seq_len, int q_dim, int half_idx) { + double total = 0.0; + for (int s = 0; s < seq_len; s++) { + const float *row_a = a + s * (q_dim * 2) + half_idx * q_dim; + const float *row_b = b + s * (q_dim * 2) + half_idx * q_dim; + for (int i = 0; i < q_dim; i++) total += fabs((double)row_a[i] - (double)row_b[i]); + } + return total / (double)(seq_len * q_dim); +} + +static double max_abs_diff_q_half(const float *a, const float *b, int seq_len, int q_dim, int half_idx) { + double best = 0.0; + for (int s = 0; s < seq_len; s++) { + const float *row_a = a + s * (q_dim * 2) + half_idx * q_dim; + const float *row_b = b + s * (q_dim * 2) + half_idx * q_dim; + for (int i = 0; i < q_dim; i++) { + double d = fabs((double)row_a[i] - (double)row_b[i]); + if (d > best) best = d; + } + } + return best; +} + +static double mean_abs_diff_q_half_swapped(const float *cpu_q, const float *ane_q, int seq_len, int q_dim) { + double total = 0.0; + for (int s = 0; s < seq_len; s++) { + const float *cpu_query = cpu_q + s * (q_dim * 2); + const float *cpu_gate = cpu_query + q_dim; + const float *ane_gate = ane_q + s * (q_dim * 2); + const float *ane_query = ane_gate + q_dim; + for (int i = 0; i < q_dim; i++) total += fabs((double)cpu_query[i] - (double)ane_query[i]); + for (int i = 0; i < q_dim; i++) total += fabs((double)cpu_gate[i] - (double)ane_gate[i]); + } + return total / (double)(seq_len * q_dim * 2); +} + +static double max_abs_diff_q_half_swapped(const float *cpu_q, const float *ane_q, int seq_len, int q_dim) { + double best = 0.0; + for (int s = 0; s < seq_len; s++) { + const float *cpu_query = cpu_q + s * (q_dim * 2); + const float *cpu_gate = cpu_query + q_dim; + const float *ane_gate = ane_q + s * (q_dim * 2); + const float *ane_query = ane_gate + q_dim; + for (int i = 0; i < q_dim; i++) { + double d = fabs((double)cpu_query[i] - (double)ane_query[i]); + if (d > best) best = d; + } + for (int i = 0; i < q_dim; i++) { + double d = fabs((double)cpu_gate[i] - (double)ane_gate[i]); + if (d > best) best = d; + } + } + return best; +} + +static double mean_abs_diff_kv_swapped(const float *cpu_k, + const float *cpu_v, + const float *ane_k, + const float *ane_v, + int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)cpu_k[i] - (double)ane_v[i]); + for (int i = 0; i < n; i++) total += fabs((double)cpu_v[i] - (double)ane_k[i]); + return total / (double)(n * 2); +} + +static double max_abs_diff_kv_swapped(const float *cpu_k, + const float *cpu_v, + const float *ane_k, + const float *ane_v, + int n) { + double best = 0.0; + for (int i = 0; i < n; i++) { + double d = fabs((double)cpu_k[i] - (double)ane_v[i]); + if (d > best) best = d; + } + for (int i = 0; i < n; i++) { + double d = fabs((double)cpu_v[i] - (double)ane_k[i]); + if (d > best) best = d; + } + return best; +} + +static inline float fp16_roundtrip(float x); + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_vproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_input_rms_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + return dict; +} + +static NSDictionary *build_vproj_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_down_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_down_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_rms_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + return dict; +} + +static NSDictionary *build_ffn_gateup_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_uponly_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_siluonly_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_down_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_down_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_down_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static int use_fp32_rrms_powchain_mode(void) { + const char *mode = getenv("ORION_RMSNORM_RRMS_MODE"); + return mode && (strcmp(mode, "nr1") == 0 || strcmp(mode, "fp32") == 0); +} + +static int qwen35_input_rmsnorm_local(OrionGraph *g, int input, int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int s = bucket; + char path[256]; + int ln_shape[4] = {1, d, 1, 1}; + snprintf(path, sizeof(path), "@model_path/layer%d/input_layernorm.bin", layer); + int rms_w = orion_gb_const_weight(g, "input_ln_w", ORION_DTYPE_FP16, ln_shape, path, 64); + return orion_gb_rmsnorm(g, input, rms_w, 1e-6f, "input_rms", d, s); +} + +static int qwen35_post_attn_rmsnorm_local(OrionGraph *g, int input, int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int s = bucket; + char path[256]; + int ln_shape[4] = {1, d, 1, 1}; + snprintf(path, sizeof(path), "@model_path/layer%d/post_attention_layernorm.bin", layer); + int rms_w = orion_gb_const_weight(g, "post_attn_ln_w", ORION_DTYPE_FP16, ln_shape, path, 64); + return orion_gb_rmsnorm(g, input, rms_w, 1e-6f, "post_attn_rms", d, s); +} + +static OrionGraph *build_ffn_rms_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int s = bucket; + int shape[4] = {1, d, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int rms = qwen35_post_attn_rmsnorm_local(g, x16, layer, bucket, cfg); + int rms32 = orion_gb_cast(g, rms, ORION_DTYPE_FP32, "rms32", shape); + orion_gb_output(g, rms32, "rms"); + return g; +} + +static OrionGraph *build_ffn_gateup_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int h = cfg->hidden_dim; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int hid_shape[4] = {1, h, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + int rms = qwen35_post_attn_rmsnorm_local(g, x16, layer, bucket, cfg); + + char gate_w[256], up_w[256]; + snprintf(gate_w, sizeof(gate_w), "@model_path/layer%d/mlp_gate_proj.bin", layer); + snprintf(up_w, sizeof(up_w), "@model_path/layer%d/mlp_up_proj.bin", layer); + + int gate = orion_gb_linear(g, rms, "gate_proj", d, h, s, gate_w, NULL); + int up = orion_gb_linear(g, rms, "up_proj", d, h, s, up_w, NULL); + int silu = orion_gb_silu(g, gate, "gate_silu", h, s); + int hidden = orion_gb_mul(g, silu, up, "ffn_hidden"); + + int gate32 = orion_gb_cast(g, gate, ORION_DTYPE_FP32, "gate32", hid_shape); + int up32 = orion_gb_cast(g, up, ORION_DTYPE_FP32, "up32", hid_shape); + int silu32 = orion_gb_cast(g, silu, ORION_DTYPE_FP32, "silu32", hid_shape); + int hidden32 = orion_gb_cast(g, hidden, ORION_DTYPE_FP32, "hidden32", hid_shape); + orion_gb_output(g, gate32, "gate_proj"); + orion_gb_output(g, up32, "up_proj"); + orion_gb_output(g, silu32, "silu_gate"); + orion_gb_output(g, hidden32, "hidden_ff"); + return g; +} + +static OrionGraph *build_ffn_uponly_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int h = cfg->hidden_dim; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int hid_shape[4] = {1, h, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + int rms = qwen35_post_attn_rmsnorm_local(g, x16, layer, bucket, cfg); + + char up_w[256]; + snprintf(up_w, sizeof(up_w), "@model_path/layer%d/mlp_up_proj.bin", layer); + int up = orion_gb_linear(g, rms, "up_proj", d, h, s, up_w, NULL); + int up32 = orion_gb_cast(g, up, ORION_DTYPE_FP32, "up32", hid_shape); + orion_gb_output(g, up32, "up_proj"); + return g; +} + +static OrionGraph *build_ffn_siluonly_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int h = cfg->hidden_dim; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int hid_shape[4] = {1, h, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + int rms = qwen35_post_attn_rmsnorm_local(g, x16, layer, bucket, cfg); + + char gate_w[256]; + snprintf(gate_w, sizeof(gate_w), "@model_path/layer%d/mlp_gate_proj.bin", layer); + int gate = orion_gb_linear(g, rms, "gate_proj", d, h, s, gate_w, NULL); + int silu = orion_gb_silu(g, gate, "gate_silu", h, s); + int silu32 = orion_gb_cast(g, silu, ORION_DTYPE_FP32, "silu32", hid_shape); + orion_gb_output(g, silu32, "silu_gate"); + return g; +} + +static OrionGraph *build_ffn_hiddenonly_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int h = cfg->hidden_dim; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int hid_shape[4] = {1, h, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + int rms = qwen35_post_attn_rmsnorm_local(g, x16, layer, bucket, cfg); + + char gate_w[256], up_w[256]; + snprintf(gate_w, sizeof(gate_w), "@model_path/layer%d/mlp_gate_proj.bin", layer); + snprintf(up_w, sizeof(up_w), "@model_path/layer%d/mlp_up_proj.bin", layer); + int gate = orion_gb_linear(g, rms, "gate_proj", d, h, s, gate_w, NULL); + int up = orion_gb_linear(g, rms, "up_proj", d, h, s, up_w, NULL); + int silu = orion_gb_silu(g, gate, "gate_silu", h, s); + int hidden = orion_gb_mul(g, silu, up, "ffn_hidden"); + int hidden32 = orion_gb_cast(g, hidden, ORION_DTYPE_FP32, "hidden32", hid_shape); + orion_gb_output(g, hidden32, "hidden_ff"); + return g; +} + +static OrionGraph *build_ffn_down_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int h = cfg->hidden_dim; + int s = bucket; + int in_shape[4] = {1, h, 1, s}; + int out_shape[4] = {1, d, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + + char down_w[256]; + snprintf(down_w, sizeof(down_w), "@model_path/layer%d/mlp_down_proj.bin", layer); + int down = orion_gb_linear(g, x16, "down_proj", h, d, s, down_w, NULL); + int down32 = orion_gb_cast(g, down, ORION_DTYPE_FP32, "down32", out_shape); + orion_gb_output(g, down32, "down_proj"); + return g; +} + +static OrionGraph *build_vproj_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int kv_heads = cfg->n_kv_head > 0 ? cfg->n_kv_head : cfg->n_head; + int kv = kv_heads * cfg->head_dim; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int out_shape[4] = {1, kv, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + int rms = qwen35_input_rmsnorm_local(g, x16, layer, bucket, cfg); + + char v_w[256]; + snprintf(v_w, sizeof(v_w), "@model_path/layer%d/self_attn_v_proj.bin", layer); + int v_proj = orion_gb_linear(g, rms, "v_proj", d, kv, s, v_w, NULL); + int v_proj32 = orion_gb_cast(g, v_proj, ORION_DTYPE_FP32, "v_proj32", out_shape); + orion_gb_output(g, v_proj32, "v_proj"); + return g; +} + +static OrionGraph *build_input_rms_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int s = bucket; + int shape[4] = {1, d, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int rms = qwen35_input_rmsnorm_local(g, x16, layer, bucket, cfg); + int rms32 = orion_gb_cast(g, rms, ORION_DTYPE_FP32, "rms32", shape); + orion_gb_output(g, rms32, "input_rms"); + return g; +} + +typedef enum { + ORION_INPUT_RMS_STAGE_MS = 0, + ORION_INPUT_RMS_STAGE_RRMS = 1, + ORION_INPUT_RMS_STAGE_XR = 2, +} OrionInputRmsStage; + +static OrionGraph *build_input_rms_stage_graph(int layer, int bucket, const OrionModelConfig *cfg, OrionInputRmsStage stage) { + int d = cfg->d_model; + int s = bucket; + int shape[4] = {1, d, 1, s}; + int scalar_shape[4] = {1, 1, 1, s}; + int ax_shape[4] = {1, 0, 0, 0}; + int axis_val = 1; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + (void)layer; + + int sq = orion_gb_mul(g, x16, x16, "input_rms_trace_sq"); + int axes = orion_gb_const_int32(g, "input_rms_trace_ax", ax_shape, &axis_val, 1); + int sum = orion_gb_reduce_sum(g, sq, axes, true, "input_rms_trace_sum", scalar_shape); + int ms = -1; + int rrms = -1; + int xr = -1; + if (use_fp32_rrms_powchain_mode()) { + int invd = orion_gb_const_scalar(g, "input_rms_trace_invd", ORION_DTYPE_FP16, 1.0f / (float)d); + ms = orion_gb_mul(g, sum, invd, "input_rms_trace_ms"); + int eps = orion_gb_const_scalar(g, "input_rms_trace_eps", ORION_DTYPE_FP16, 1e-6f); + int mse = orion_gb_add(g, ms, eps, "input_rms_trace_mse"); + int nhalf = orion_gb_const_scalar(g, "input_rms_trace_nhalf", ORION_DTYPE_FP16, -0.5f); + int rrms0 = orion_gb_pow(g, mse, nhalf, "input_rms_trace_rrms"); + + int rrms_sq = orion_gb_mul(g, rrms0, rrms0, "input_rms_trace_rrms_sq"); + int nr_term = orion_gb_mul(g, mse, rrms_sq, "input_rms_trace_nr_term"); + int half_nr = orion_gb_const_scalar(g, "input_rms_trace_half_nr", ORION_DTYPE_FP16, 0.5f); + int nr_half = orion_gb_mul(g, nr_term, half_nr, "input_rms_trace_nr_half"); + int threehalves = orion_gb_const_scalar(g, "input_rms_trace_threehalves", ORION_DTYPE_FP16, 1.5f); + int nr_corr = orion_gb_sub(g, threehalves, nr_half, "input_rms_trace_nr_corr"); + rrms = orion_gb_mul(g, rrms0, nr_corr, "input_rms_trace_rrms_refined"); + xr = orion_gb_mul(g, x16, rrms, "input_rms_trace_xr"); + } else { + int invd = orion_gb_const_scalar(g, "input_rms_trace_invd", ORION_DTYPE_FP16, 1.0f / (float)d); + ms = orion_gb_mul(g, sum, invd, "input_rms_trace_ms"); + int eps = orion_gb_const_scalar(g, "input_rms_trace_eps", ORION_DTYPE_FP16, 1e-6f); + int mse = orion_gb_add(g, ms, eps, "input_rms_trace_mse"); + int nhalf = orion_gb_const_scalar(g, "input_rms_trace_nhalf", ORION_DTYPE_FP16, -0.5f); + rrms = orion_gb_pow(g, mse, nhalf, "input_rms_trace_rrms"); + xr = orion_gb_mul(g, x16, rrms, "input_rms_trace_xr"); + } + switch (stage) { + case ORION_INPUT_RMS_STAGE_MS: { + int ms32 = orion_gb_cast(g, ms, ORION_DTYPE_FP32, "input_rms_trace_ms_out32", scalar_shape); + orion_gb_output(g, ms32, "input_rms_ms"); + break; + } + case ORION_INPUT_RMS_STAGE_RRMS: { + int rrms32 = orion_gb_cast(g, rrms, ORION_DTYPE_FP32, "input_rms_trace_rrms_out32", scalar_shape); + orion_gb_output(g, rrms32, "input_rms_rrms"); + break; + } + case ORION_INPUT_RMS_STAGE_XR: { + int xr32 = orion_gb_cast(g, xr, ORION_DTYPE_FP32, "input_rms_trace_xr_out32", shape); + orion_gb_output(g, xr32, "input_rms_xr"); + break; + } + } + return g; +} + +static OrionGraph *build_vproj_linear_only_graph(int layer, int bucket, const OrionModelConfig *cfg) { + int d = cfg->d_model; + int kv_heads = cfg->n_kv_head > 0 ? cfg->n_kv_head : cfg->n_head; + int kv = kv_heads * cfg->head_dim; + int s = bucket; + int in_shape[4] = {1, d, 1, s}; + int out_shape[4] = {1, kv, 1, s}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, in_shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", in_shape); + + char v_w[256]; + snprintf(v_w, sizeof(v_w), "@model_path/layer%d/self_attn_v_proj.bin", layer); + int v_proj = orion_gb_linear(g, x16, "v_proj_linear_only", d, kv, s, v_w, NULL); + int v_proj32 = orion_gb_cast(g, v_proj, ORION_DTYPE_FP32, "v_proj32", out_shape); + orion_gb_output(g, v_proj32, "v_proj_linear_only"); + return g; +} + +static OrionGraph *build_silu_micro_graph(int dim, int bucket) { + int shape[4] = {1, dim, 1, bucket}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int silu = orion_gb_silu(g, x16, "micro_silu", dim, bucket); + int silu32 = orion_gb_cast(g, silu, ORION_DTYPE_FP32, "silu32", shape); + orion_gb_output(g, silu32, "silu"); + return g; +} + +static OrionGraph *build_sigmoid_micro_graph(int dim, int bucket) { + int shape[4] = {1, dim, 1, bucket}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int sig = orion_gb_sigmoid(g, x16, "micro_sigmoid"); + int sig32 = orion_gb_cast(g, sig, ORION_DTYPE_FP32, "sig32", shape); + orion_gb_output(g, sig32, "sigmoid"); + return g; +} + +static int build_sigmoid_tanh_node(OrionGraph *g, int input, const char *prefix) { + char buf[ORION_MAX_NAME]; + + snprintf(buf, sizeof(buf), "%s_half", prefix); + int half = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, 0.5f); + snprintf(buf, sizeof(buf), "%s_hx", prefix); + int hx = orion_gb_mul(g, input, half, buf); + + snprintf(buf, sizeof(buf), "%s_th", prefix); + int th = orion_gb_tanh(g, hx, buf); + + snprintf(buf, sizeof(buf), "%s_one", prefix); + int one = orion_gb_const_scalar(g, buf, ORION_DTYPE_FP16, 1.0f); + snprintf(buf, sizeof(buf), "%s_onep", prefix); + int onep = orion_gb_add(g, th, one, buf); + + snprintf(buf, sizeof(buf), "%s_out", prefix); + return orion_gb_mul(g, onep, half, buf); +} + +static OrionGraph *build_sigmoid_tanh_micro_graph(int dim, int bucket) { + int shape[4] = {1, dim, 1, bucket}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int sig = build_sigmoid_tanh_node(g, x16, "micro_sigmoid_tanh"); + int sig32 = orion_gb_cast(g, sig, ORION_DTYPE_FP32, "sig32", shape); + orion_gb_output(g, sig32, "sigmoid_tanh"); + return g; +} + +static OrionGraph *build_silu_tanh_micro_graph(int dim, int bucket) { + int shape[4] = {1, dim, 1, bucket}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int sig = build_sigmoid_tanh_node(g, x16, "micro_silu_tanh_sig"); + int silu = orion_gb_mul(g, x16, sig, "micro_silu_tanh"); + int silu32 = orion_gb_cast(g, silu, ORION_DTYPE_FP32, "silu32", shape); + orion_gb_output(g, silu32, "silu_tanh"); + return g; +} + +static OrionGraph *build_mul_micro_graph(int dim, int bucket) { + int shape[4] = {1, dim, 1, bucket}; + OrionGraph *g = orion_graph_create(); + int a = orion_gb_input(g, "a", ORION_DTYPE_FP32, shape); + int b = orion_gb_input(g, "b", ORION_DTYPE_FP32, shape); + int a16 = orion_gb_cast(g, a, ORION_DTYPE_FP16, "a16", shape); + int b16 = orion_gb_cast(g, b, ORION_DTYPE_FP16, "b16", shape); + int mul = orion_gb_mul(g, a16, b16, "micro_mul"); + int mul32 = orion_gb_cast(g, mul, ORION_DTYPE_FP32, "mul32", shape); + orion_gb_output(g, mul32, "mul"); + return g; +} + +static OrionGraph *build_rrms_pow_micro_graph(int bucket) { + int shape[4] = {1, 1, 1, bucket}; + OrionGraph *g = orion_graph_create(); + int x = orion_gb_input(g, "x", ORION_DTYPE_FP32, shape); + int x16 = orion_gb_cast(g, x, ORION_DTYPE_FP16, "x16", shape); + int nhalf = orion_gb_const_scalar(g, "rrms_pow_nhalf", ORION_DTYPE_FP16, -0.5f); + int out = orion_gb_pow(g, x16, nhalf, "rrms_pow"); + int out32 = orion_gb_cast(g, out, ORION_DTYPE_FP32, "rrms_pow32", shape); + orion_gb_output(g, out32, "rrms_pow"); + return g; +} + +typedef struct { + OrionProgram *prog_q; + OrionProgram *prog_kv; + OrionProgram *prog_v_only; + OrionProgram *prog_input_rms_only; + OrionProgram *prog_input_rms_ms; + OrionProgram *prog_input_rms_rrms; + OrionProgram *prog_input_rms_xr; + OrionProgram *prog_v_linear_only; + OrionProgram *prog_ffn_rms; + OrionProgram *prog_ffn_gateup; + OrionProgram *prog_ffn_up_only; + OrionProgram *prog_ffn_silu_only; + OrionProgram *prog_ffn_hidden_only; + OrionProgram *prog_ffn_down; + OrionProgram *prog_ffn; + OrionProgram *prog_sigmoid_micro; + OrionProgram *prog_sigmoid_tanh_micro; + OrionProgram *prog_silu_micro; + OrionProgram *prog_silu_tanh_micro; + OrionProgram *prog_mul_micro; + OrionProgram *prog_rrms_pow_micro; + int layer_idx; + int bucket; + int d_model; + int q_dim; + int kv_dim; +} OrionQwen35AneBridge; + +static void bridge_release(OrionQwen35AneBridge *bridge) { + if (!bridge) return; + if (bridge->prog_q) orion_release_program(bridge->prog_q); + if (bridge->prog_kv) orion_release_program(bridge->prog_kv); + if (bridge->prog_v_only) orion_release_program(bridge->prog_v_only); + if (bridge->prog_input_rms_only) orion_release_program(bridge->prog_input_rms_only); + if (bridge->prog_input_rms_ms) orion_release_program(bridge->prog_input_rms_ms); + if (bridge->prog_input_rms_rrms) orion_release_program(bridge->prog_input_rms_rrms); + if (bridge->prog_input_rms_xr) orion_release_program(bridge->prog_input_rms_xr); + if (bridge->prog_v_linear_only) orion_release_program(bridge->prog_v_linear_only); + if (bridge->prog_ffn_rms) orion_release_program(bridge->prog_ffn_rms); + if (bridge->prog_ffn_gateup) orion_release_program(bridge->prog_ffn_gateup); + if (bridge->prog_ffn_up_only) orion_release_program(bridge->prog_ffn_up_only); + if (bridge->prog_ffn_silu_only) orion_release_program(bridge->prog_ffn_silu_only); + if (bridge->prog_ffn_hidden_only) orion_release_program(bridge->prog_ffn_hidden_only); + if (bridge->prog_ffn_down) orion_release_program(bridge->prog_ffn_down); + if (bridge->prog_ffn) orion_release_program(bridge->prog_ffn); + if (bridge->prog_sigmoid_micro) orion_release_program(bridge->prog_sigmoid_micro); + if (bridge->prog_sigmoid_tanh_micro) orion_release_program(bridge->prog_sigmoid_tanh_micro); + if (bridge->prog_silu_micro) orion_release_program(bridge->prog_silu_micro); + if (bridge->prog_silu_tanh_micro) orion_release_program(bridge->prog_silu_tanh_micro); + if (bridge->prog_mul_micro) orion_release_program(bridge->prog_mul_micro); + if (bridge->prog_rrms_pow_micro) orion_release_program(bridge->prog_rrms_pow_micro); + memset(bridge, 0, sizeof(*bridge)); +} + +static int bridge_init(OrionQwen35AneBridge *bridge, NSString *blobDir, int layer, int bucket, OrionQwen35Manifest *manifest) { + memset(bridge, 0, sizeof(*bridge)); + OrionModelConfig cfg = { + .n_layer = manifest->n_layer, + .n_head = manifest->n_head, + .n_kv_head = manifest->n_kv_head, + .d_model = manifest->d_model, + .head_dim = manifest->head_dim, + .hidden_dim = manifest->d_ff, + .vocab = manifest->vocab, + .max_seq = manifest->max_seq, + }; + NSString *mil_q = compile_graph(orion_frontend_qwen35_prefill_q_proj(layer, bucket, &cfg)); + NSString *mil_kv = compile_graph(orion_frontend_qwen35_prefill_kv_proj(layer, bucket, &cfg)); + NSString *mil_v_only = compile_graph(build_vproj_graph(layer, bucket, &cfg)); + NSString *mil_input_rms_only = compile_graph(build_input_rms_graph(layer, bucket, &cfg)); + NSString *mil_input_rms_ms = compile_graph(build_input_rms_stage_graph(layer, bucket, &cfg, ORION_INPUT_RMS_STAGE_MS)); + NSString *mil_input_rms_rrms = compile_graph(build_input_rms_stage_graph(layer, bucket, &cfg, ORION_INPUT_RMS_STAGE_RRMS)); + NSString *mil_input_rms_xr = compile_graph(build_input_rms_stage_graph(layer, bucket, &cfg, ORION_INPUT_RMS_STAGE_XR)); + NSString *mil_v_linear_only = compile_graph(build_vproj_linear_only_graph(layer, bucket, &cfg)); + NSString *mil_ffn_rms = compile_graph(build_ffn_rms_graph(layer, bucket, &cfg)); + NSString *mil_ffn_gateup = compile_graph(build_ffn_gateup_graph(layer, bucket, &cfg)); + NSString *mil_ffn_up_only = compile_graph(build_ffn_uponly_graph(layer, bucket, &cfg)); + NSString *mil_ffn_silu_only = compile_graph(build_ffn_siluonly_graph(layer, bucket, &cfg)); + NSString *mil_ffn_hidden_only = compile_graph(build_ffn_hiddenonly_graph(layer, bucket, &cfg)); + NSString *mil_ffn_down = compile_graph(build_ffn_down_graph(layer, bucket, &cfg)); + NSString *mil_ffn = compile_graph(orion_frontend_qwen35_prefill_ffn(layer, bucket, &cfg)); + NSString *mil_sigmoid_micro = compile_graph(build_sigmoid_micro_graph(cfg.hidden_dim, bucket)); + NSString *mil_sigmoid_tanh_micro = compile_graph(build_sigmoid_tanh_micro_graph(cfg.hidden_dim, bucket)); + NSString *mil_silu_micro = compile_graph(build_silu_micro_graph(cfg.hidden_dim, bucket)); + NSString *mil_silu_tanh_micro = compile_graph(build_silu_tanh_micro_graph(cfg.hidden_dim, bucket)); + NSString *mil_mul_micro = compile_graph(build_mul_micro_graph(cfg.hidden_dim, bucket)); + NSString *mil_rrms_pow_micro = compile_graph(build_rrms_pow_micro_graph(bucket)); + if (!mil_q || !mil_kv || !mil_v_only || !mil_input_rms_only || !mil_input_rms_ms || !mil_input_rms_rrms || + !mil_input_rms_xr || !mil_v_linear_only || + !mil_ffn_rms || !mil_ffn_gateup || !mil_ffn_up_only || + !mil_ffn_silu_only || !mil_ffn_hidden_only || !mil_ffn_down || !mil_ffn || + !mil_sigmoid_micro || !mil_sigmoid_tanh_micro || !mil_silu_micro || !mil_silu_tanh_micro || + !mil_mul_micro || !mil_rrms_pow_micro) return 0; + + bridge->prog_q = orion_compile_mil(mil_q.UTF8String, build_qproj_wdict(layer, blobDir), "qwen35_9b_stage_q"); + bridge->prog_kv = orion_compile_mil(mil_kv.UTF8String, build_kv_wdict(layer, blobDir), "qwen35_9b_stage_kv"); + bridge->prog_v_only = orion_compile_mil(mil_v_only.UTF8String, build_vproj_wdict(layer, blobDir), "qwen35_9b_stage_v_only"); + bridge->prog_input_rms_only = orion_compile_mil(mil_input_rms_only.UTF8String, build_input_rms_wdict(layer, blobDir), "qwen35_9b_stage_input_rms_only"); + bridge->prog_input_rms_ms = orion_compile_mil(mil_input_rms_ms.UTF8String, build_input_rms_wdict(layer, blobDir), "qwen35_9b_stage_input_rms_ms"); + bridge->prog_input_rms_rrms = orion_compile_mil(mil_input_rms_rrms.UTF8String, build_input_rms_wdict(layer, blobDir), "qwen35_9b_stage_input_rms_rrms"); + bridge->prog_input_rms_xr = orion_compile_mil(mil_input_rms_xr.UTF8String, build_input_rms_wdict(layer, blobDir), "qwen35_9b_stage_input_rms_xr"); + bridge->prog_v_linear_only = orion_compile_mil(mil_v_linear_only.UTF8String, build_vproj_linear_only_wdict(layer, blobDir), "qwen35_9b_stage_v_linear_only"); + bridge->prog_ffn_rms = orion_compile_mil(mil_ffn_rms.UTF8String, build_ffn_rms_wdict(layer, blobDir), "qwen35_9b_stage_ffn_rms"); + bridge->prog_ffn_gateup = orion_compile_mil(mil_ffn_gateup.UTF8String, build_ffn_gateup_wdict(layer, blobDir), "qwen35_9b_stage_ffn_gateup"); + bridge->prog_ffn_up_only = orion_compile_mil(mil_ffn_up_only.UTF8String, build_ffn_uponly_wdict(layer, blobDir), "qwen35_9b_stage_ffn_up_only"); + bridge->prog_ffn_silu_only = orion_compile_mil(mil_ffn_silu_only.UTF8String, build_ffn_siluonly_wdict(layer, blobDir), "qwen35_9b_stage_ffn_silu_only"); + bridge->prog_ffn_hidden_only = orion_compile_mil(mil_ffn_hidden_only.UTF8String, build_ffn_gateup_wdict(layer, blobDir), "qwen35_9b_stage_ffn_hidden_only"); + bridge->prog_ffn_down = orion_compile_mil(mil_ffn_down.UTF8String, build_ffn_down_wdict(layer, blobDir), "qwen35_9b_stage_ffn_down"); + bridge->prog_ffn = orion_compile_mil(mil_ffn.UTF8String, build_ffn_wdict(layer, blobDir), "qwen35_9b_stage_ffn"); + bridge->prog_sigmoid_micro = orion_compile_mil(mil_sigmoid_micro.UTF8String, nil, "qwen35_9b_stage_sigmoid_micro"); + bridge->prog_sigmoid_tanh_micro = orion_compile_mil(mil_sigmoid_tanh_micro.UTF8String, nil, "qwen35_9b_stage_sigmoid_tanh_micro"); + bridge->prog_silu_micro = orion_compile_mil(mil_silu_micro.UTF8String, nil, "qwen35_9b_stage_silu_micro"); + bridge->prog_silu_tanh_micro = orion_compile_mil(mil_silu_tanh_micro.UTF8String, nil, "qwen35_9b_stage_silu_tanh_micro"); + bridge->prog_mul_micro = orion_compile_mil(mil_mul_micro.UTF8String, nil, "qwen35_9b_stage_mul_micro"); + bridge->prog_rrms_pow_micro = orion_compile_mil(mil_rrms_pow_micro.UTF8String, nil, "qwen35_9b_stage_rrms_pow_micro"); + if (!bridge->prog_q || !bridge->prog_kv || !bridge->prog_v_only || + !bridge->prog_input_rms_only || !bridge->prog_input_rms_ms || !bridge->prog_input_rms_rrms || + !bridge->prog_input_rms_xr || !bridge->prog_v_linear_only || !bridge->prog_ffn_rms || + !bridge->prog_ffn_gateup || !bridge->prog_ffn_up_only || !bridge->prog_ffn_silu_only || + !bridge->prog_ffn_hidden_only || + !bridge->prog_ffn_down || !bridge->prog_ffn || !bridge->prog_sigmoid_micro || + !bridge->prog_sigmoid_tanh_micro || !bridge->prog_silu_micro || !bridge->prog_silu_tanh_micro || + !bridge->prog_mul_micro || !bridge->prog_rrms_pow_micro) { + bridge_release(bridge); + return 0; + } + + bridge->layer_idx = layer; + bridge->bucket = bucket; + bridge->d_model = manifest->d_model; + bridge->q_dim = manifest->n_head * manifest->head_dim; + bridge->kv_dim = manifest->n_kv_head * manifest->head_dim; + return 1; +} + +static int load_embeddings(const char *blob_dir, + OrionQwen35Manifest *manifest, + const int *token_ids, + int seq_len, + float *hidden_out) { + const int d_model = manifest->d_model; + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + for (int s = 0; s < seq_len; s++) { + if (!orion_read_blob_row_f32(embed_path, token_ids[s], d_model, hidden_out + s * d_model)) return 0; + } + return 1; +} + +static int apply_cpu_ffn(const float *hidden_attn, + int seq_len, + int d_model, + int d_ff, + const float *post_ln, + const float *gate_proj, + const float *up_proj, + const float *down_proj, + float *hidden_final_out) { + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!scratch || !mlp_out) { + free(scratch); + free(mlp_out); + return 0; + } + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_final_out[i] = hidden_attn[i] + mlp_out[i]; + free(scratch); + free(mlp_out); + return 1; +} + +static int apply_cpu_layer(const char *blob_dir, + OrionQwen35Manifest *manifest, + int layer_idx, + const float *hidden_in, + int seq_len, + float *hidden_out) { + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + return 0; + } + + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *attn_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!normed || !attn_out || !hidden_attn) { + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(attn_out); free(hidden_attn); + return 0; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_in + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + if (file_exists(full_q_path)) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(attn_out); free(hidden_attn); + return 0; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + char path_qkv[2048], path_out[2048], path_dt[2048], path_norm[2048], path_conv[2048]; + snprintf(path_qkv, sizeof(path_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + snprintf(path_out, sizeof(path_out), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx); + snprintf(path_dt, sizeof(path_dt), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + snprintf(path_norm, sizeof(path_norm), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx); + snprintf(path_conv, sizeof(path_conv), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + + int qkv_rows = orion_blob_element_count(path_qkv) / d_model; + int value_dim = orion_blob_element_count(path_out) / d_model; + int num_v_heads = orion_blob_element_count(path_dt); + int head_v_dim = orion_blob_element_count(path_norm); + int key_dim = (qkv_rows - value_dim) / 2; + int num_k_heads = num_v_heads; + int head_k_dim = key_dim / num_k_heads; + int conv_kernel = orion_blob_element_count(path_conv) / qkv_rows; + + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(attn_out); free(hidden_attn); + return 0; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + attn_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) hidden_attn[i] = hidden_in[i] + attn_out[i]; + int ok = apply_cpu_ffn(hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, hidden_out); + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(attn_out); free(hidden_attn); + return ok; +} + +static void cpu_linear_batch(const float *x_seq, + int seq_len, + const float *weight, + int in_dim, + int out_dim, + float *out_seq) { + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, out_dim, in_dim, + 1.0f, x_seq, in_dim, weight, in_dim, + 0.0f, out_seq, out_dim); +} + +static void cpu_linear_batch_fp16_emulated(const float *x_seq, + int seq_len, + const float *weight, + int in_dim, + int out_dim, + float *out_seq) { + for (int s = 0; s < seq_len; s++) { + const float *x = x_seq + (size_t)s * in_dim; + float *out = out_seq + (size_t)s * out_dim; + for (int o = 0; o < out_dim; o++) { + const float *w = weight + (size_t)o * in_dim; + float acc = 0.0f; + for (int i = 0; i < in_dim; i++) { + float x16 = fp16_roundtrip(x[i]); + float w16 = fp16_roundtrip(w[i]); + float prod = fp16_roundtrip(x16 * w16); + acc = fp16_roundtrip(acc + prod); + } + out[o] = acc; + } + } +} + +static void cpu_linear_batch_fp16in_fp32acc_fp16out(const float *x_seq, + int seq_len, + const float *weight, + int in_dim, + int out_dim, + float *out_seq) { + for (int s = 0; s < seq_len; s++) { + const float *x = x_seq + (size_t)s * in_dim; + float *out = out_seq + (size_t)s * out_dim; + for (int o = 0; o < out_dim; o++) { + const float *w = weight + (size_t)o * in_dim; + float acc = 0.0f; + for (int i = 0; i < in_dim; i++) { + float x16 = fp16_roundtrip(x[i]); + float w16 = fp16_roundtrip(w[i]); + acc += x16 * w16; + } + out[o] = fp16_roundtrip(acc); + } + } +} + +static int selected_token_logits(const char *blob_dir, + const char *lm_head_name, + const float *hidden, + int d_model, + const int *token_ids, + int token_count, + float *out_logits) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + for (int tok_idx = 0; tok_idx < token_count; tok_idx++) { + int tok = token_ids[tok_idx]; + if (tok < 0 || !orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + out_logits[tok_idx] = dot; + } + + free(row); + return 1; +} + +#define ORION_TRACE_DIM_MAX 16 +#define ORION_TRACE_ATTR_TOPK 4 + +typedef struct { + int channel; + double abs_contrib; + double contrib; + double delta_input; + double weight; +} OrionTraceChannelImpact; + +static int compare_trace_channel_impact_desc(const void *lhs, const void *rhs) { + const OrionTraceChannelImpact *a = (const OrionTraceChannelImpact *)lhs; + const OrionTraceChannelImpact *b = (const OrionTraceChannelImpact *)rhs; + if (a->abs_contrib > b->abs_contrib) return -1; + if (a->abs_contrib < b->abs_contrib) return 1; + if (a->channel < b->channel) return -1; + if (a->channel > b->channel) return 1; + return 0; +} + +static int parse_dim_list(const char *csv, int *out_dims, int max_dims, int dim_limit) { + if (!csv || !out_dims || max_dims <= 0) return 0; + char *copy = strdup(csv); + if (!copy) return 0; + int count = 0; + char *save = NULL; + for (char *tok = strtok_r(copy, ",", &save); tok && count < max_dims; tok = strtok_r(NULL, ",", &save)) { + while (*tok && isspace((unsigned char)*tok)) tok++; + if (!*tok) continue; + char *end = NULL; + long dim = strtol(tok, &end, 10); + if (end == tok) continue; + while (*end && isspace((unsigned char)*end)) end++; + if (*end != '\0') continue; + if (dim < 0 || dim >= dim_limit) continue; + int duplicate = 0; + for (int i = 0; i < count; i++) { + if (out_dims[i] == (int)dim) { + duplicate = 1; + break; + } + } + if (duplicate) continue; + out_dims[count++] = (int)dim; + } + free(copy); + return count; +} + +static int use_cpu_v_proj_override(void) { + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_proj_override(void) { + const char *source = getenv("ORION_Q_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_query_override(void) { + const char *source = getenv("ORION_Q_QUERY_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_gate_override(void) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int load_v_proj_cpu_channel_override(int *out_channels, int max_channels, int q_dim) { + const char *csv = getenv("ORION_V_PROJ_CPU_CHANNELS"); + if (!csv || !*csv) return 0; + return parse_dim_list(csv, out_channels, max_channels, q_dim); +} + +static const char *q_proj_source_label(int q_proj_uses_cpu) { + return q_proj_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *q_query_source_label(int q_query_uses_cpu) { + return q_query_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *q_gate_source_label(int q_gate_uses_cpu) { + return q_gate_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *v_proj_source_label(int v_proj_uses_cpu, int partial_channel_count) { + if (v_proj_uses_cpu) return "cpu_linear"; + if (partial_channel_count > 0) return "ane_linear_cpu_channels"; + return "ane_linear"; +} + +static void apply_v_proj_cpu_channel_overrides(float *dst_v_proj_seq, + const float *cpu_v_proj_seq, + int seq_len, + int kv_dim, + int head_dim, + int q_per_kv, + const int *channels, + int channel_count) { + if (!dst_v_proj_seq || !cpu_v_proj_seq || !channels || channel_count <= 0 || seq_len <= 0 || + kv_dim <= 0 || head_dim <= 0 || q_per_kv <= 0) { + return; + } + for (int i = 0; i < channel_count; i++) { + int channel = channels[i]; + int head = channel / head_dim; + int offset = channel % head_dim; + int kv_head = head / q_per_kv; + int kv_channel = kv_head * head_dim + offset; + if (kv_channel < 0 || kv_channel >= kv_dim) continue; + for (int s = 0; s < seq_len; s++) { + dst_v_proj_seq[(size_t)s * kv_dim + kv_channel] = + cpu_v_proj_seq[(size_t)s * kv_dim + kv_channel]; + } + } +} + +static void print_trace_stage_rows(const char *stage, + const float *cpu_seq, + const float *hybrid_seq, + int seq_len, + int d_model, + const int *trace_dims, + int trace_dim_count) { + if (!stage || !cpu_seq || !hybrid_seq || !trace_dims || trace_dim_count <= 0 || seq_len <= 0) return; + const float *cpu_last = cpu_seq + (seq_len - 1) * d_model; + const float *hybrid_last = hybrid_seq + (seq_len - 1) * d_model; + for (int i = 0; i < trace_dim_count; i++) { + int dim = trace_dims[i]; + float cpu_val = cpu_last[dim]; + float hybrid_val = hybrid_last[dim]; + printf("trace_stage_dim stage=%s dim=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + stage, dim, cpu_val, hybrid_val, fabsf(cpu_val - hybrid_val)); + } +} + +static void print_trace_attn_attr_rows(const float *cpu_in_seq, + const float *hybrid_in_seq, + const float *o_proj, + int seq_len, + int q_dim, + int head_dim, + const int *trace_dims, + int trace_dim_count) { + if (!cpu_in_seq || !hybrid_in_seq || !o_proj || !trace_dims || trace_dim_count <= 0 || seq_len <= 0) return; + const float *cpu_last = cpu_in_seq + (seq_len - 1) * q_dim; + const float *hybrid_last = hybrid_in_seq + (seq_len - 1) * q_dim; + OrionTraceChannelImpact top[ORION_TRACE_ATTR_TOPK]; + for (int dim_idx = 0; dim_idx < trace_dim_count; dim_idx++) { + int out_dim = trace_dims[dim_idx]; + int top_count = 0; + const float *weight_row = o_proj + (size_t)out_dim * q_dim; + for (int channel = 0; channel < q_dim; channel++) { + double delta_input = (double)hybrid_last[channel] - (double)cpu_last[channel]; + double weight = (double)weight_row[channel]; + double contrib = delta_input * weight; + OrionTraceChannelImpact cand = { + .channel = channel, + .abs_contrib = fabs(contrib), + .contrib = contrib, + .delta_input = delta_input, + .weight = weight, + }; + if (top_count < ORION_TRACE_ATTR_TOPK) { + top[top_count++] = cand; + continue; + } + int worst_idx = 0; + for (int i = 1; i < top_count; i++) { + if (top[i].abs_contrib < top[worst_idx].abs_contrib) worst_idx = i; + } + if (cand.abs_contrib > top[worst_idx].abs_contrib) top[worst_idx] = cand; + } + qsort(top, (size_t)top_count, sizeof(top[0]), compare_trace_channel_impact_desc); + for (int rank = 0; rank < top_count; rank++) { + int channel = top[rank].channel; + int head = (head_dim > 0) ? (channel / head_dim) : -1; + int offset = (head_dim > 0) ? (channel % head_dim) : channel; + printf("trace_attn_attr dim=%d rank=%d channel=%d head=%d offset=%d delta_input=%.6f weight=%.6f contrib=%.6f\n", + out_dim, rank + 1, channel, head, offset, + top[rank].delta_input, top[rank].weight, top[rank].contrib); + } + } +} + +static void print_trace_attn_channel_rows(const float *cpu_q_proj_seq, + const float *hybrid_q_proj_seq, + const float *cpu_v_proj_seq, + const float *cpu_v_rms_only_proj_seq, + const float *hybrid_v_proj_seq, + const float *hybrid_v_single_proj_seq, + const float *hybrid_v_linear_only_proj_seq, + const float *cpu_q_normed_seq, + const float *hybrid_q_normed_seq, + const float *cpu_gate_sigmoid_seq, + const float *hybrid_gate_sigmoid_seq, + const float *cpu_q_rope_seq, + const float *hybrid_q_rope_seq, + const float *cpu_context_seq, + const float *hybrid_context_seq, + const float *cpu_gated_context_seq, + const float *hybrid_gated_context_seq, + int seq_len, + int q_dim, + int kv_dim, + int head_dim, + int q_per_kv, + const int *trace_channels, + int trace_channel_count) { + if (!cpu_q_proj_seq || !hybrid_q_proj_seq || !cpu_v_proj_seq || !cpu_v_rms_only_proj_seq || + !hybrid_v_proj_seq || + !hybrid_v_single_proj_seq || !hybrid_v_linear_only_proj_seq || + !cpu_q_normed_seq || !hybrid_q_normed_seq || !cpu_gate_sigmoid_seq || !hybrid_gate_sigmoid_seq || + !cpu_q_rope_seq || !hybrid_q_rope_seq || !cpu_context_seq || !hybrid_context_seq || + !cpu_gated_context_seq || !hybrid_gated_context_seq || !trace_channels || + trace_channel_count <= 0 || seq_len <= 0 || head_dim <= 0 || q_per_kv <= 0) { + return; + } + + const float *cpu_q_last = cpu_q_proj_seq + (seq_len - 1) * (q_dim * 2); + const float *hybrid_q_last = hybrid_q_proj_seq + (seq_len - 1) * (q_dim * 2); + const float *cpu_v_last = cpu_v_proj_seq + (seq_len - 1) * kv_dim; + const float *cpu_v_rms_only_last = cpu_v_rms_only_proj_seq + (seq_len - 1) * kv_dim; + const float *hybrid_v_last = hybrid_v_proj_seq + (seq_len - 1) * kv_dim; + const float *hybrid_v_single_last = hybrid_v_single_proj_seq + (seq_len - 1) * kv_dim; + const float *hybrid_v_linear_only_last = hybrid_v_linear_only_proj_seq + (seq_len - 1) * kv_dim; + const float *cpu_q_norm_last = cpu_q_normed_seq + (seq_len - 1) * q_dim; + const float *hybrid_q_norm_last = hybrid_q_normed_seq + (seq_len - 1) * q_dim; + const float *cpu_gate_sigmoid_last = cpu_gate_sigmoid_seq + (seq_len - 1) * q_dim; + const float *hybrid_gate_sigmoid_last = hybrid_gate_sigmoid_seq + (seq_len - 1) * q_dim; + const float *cpu_q_rope_last = cpu_q_rope_seq + (seq_len - 1) * q_dim; + const float *hybrid_q_rope_last = hybrid_q_rope_seq + (seq_len - 1) * q_dim; + const float *cpu_context_last = cpu_context_seq + (seq_len - 1) * q_dim; + const float *hybrid_context_last = hybrid_context_seq + (seq_len - 1) * q_dim; + const float *cpu_gated_context_last = cpu_gated_context_seq + (seq_len - 1) * q_dim; + const float *hybrid_gated_context_last = hybrid_gated_context_seq + (seq_len - 1) * q_dim; + + for (int i = 0; i < trace_channel_count; i++) { + int channel = trace_channels[i]; + int head = channel / head_dim; + int offset = channel % head_dim; + int kv_head = head / q_per_kv; + int kv_channel = kv_head * head_dim + offset; + if (kv_channel < 0 || kv_channel >= kv_dim) continue; + + float cpu_query = cpu_q_last[channel]; + float hybrid_query = hybrid_q_last[channel]; + float cpu_gate = cpu_q_last[q_dim + channel]; + float hybrid_gate = hybrid_q_last[q_dim + channel]; + float cpu_q_norm = cpu_q_norm_last[channel]; + float hybrid_q_norm = hybrid_q_norm_last[channel]; + float cpu_gate_sig = cpu_gate_sigmoid_last[channel]; + float hybrid_gate_sig = hybrid_gate_sigmoid_last[channel]; + float cpu_q_rope = cpu_q_rope_last[channel]; + float hybrid_q_rope = hybrid_q_rope_last[channel]; + float cpu_v = cpu_v_last[kv_channel]; + float hybrid_v = hybrid_v_last[kv_channel]; + float hybrid_v_single = hybrid_v_single_last[kv_channel]; + float cpu_context = cpu_context_last[channel]; + float hybrid_context = hybrid_context_last[channel]; + float cpu_gated_context = cpu_gated_context_last[channel]; + float hybrid_gated_context = hybrid_gated_context_last[channel]; + + printf("trace_attn_channel stage=q_query_raw channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_query, hybrid_query, fabsf(cpu_query - hybrid_query)); + printf("trace_attn_channel stage=q_gate_raw channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_gate, hybrid_gate, fabsf(cpu_gate - hybrid_gate)); + printf("trace_attn_channel stage=q_normed channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_q_norm, hybrid_q_norm, fabsf(cpu_q_norm - hybrid_q_norm)); + printf("trace_attn_channel stage=gate_sigmoid channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_gate_sig, hybrid_gate_sig, fabsf(cpu_gate_sig - hybrid_gate_sig)); + printf("trace_attn_channel stage=q_rope channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_q_rope, hybrid_q_rope, fabsf(cpu_q_rope - hybrid_q_rope)); + printf("trace_attn_channel stage=v_proj_mapped channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_v, hybrid_v, fabsf(cpu_v - hybrid_v)); + printf("trace_attn_channel stage=v_proj_rms_only_mapped channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_v, cpu_v_rms_only_last[kv_channel], + fabsf(cpu_v - cpu_v_rms_only_last[kv_channel])); + printf("trace_attn_channel stage=v_proj_single_mapped channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_v, hybrid_v_single, fabsf(cpu_v - hybrid_v_single)); + printf("trace_attn_channel stage=v_proj_linear_only_mapped channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_v, hybrid_v_linear_only_last[kv_channel], + fabsf(cpu_v - hybrid_v_linear_only_last[kv_channel])); + printf("trace_attn_channel stage=context channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_context, hybrid_context, fabsf(cpu_context - hybrid_context)); + printf("trace_attn_channel stage=gated_context channel=%d head=%d offset=%d kv_head=%d kv_channel=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + channel, head, offset, kv_head, kv_channel, cpu_gated_context, hybrid_gated_context, fabsf(cpu_gated_context - hybrid_gated_context)); + } +} + +static void print_trace_v_neighbor_rows(const char *source, + const float *cpu_v_proj_seq, + const float *hybrid_v_proj_seq, + int seq_len, + int kv_dim, + int head_dim, + int q_per_kv, + const int *trace_channels, + int trace_channel_count) { + if (!source || !cpu_v_proj_seq || !hybrid_v_proj_seq || !trace_channels || + trace_channel_count <= 0 || seq_len <= 0 || kv_dim <= 0 || head_dim <= 0 || q_per_kv <= 0) { + return; + } + + const float *cpu_v_last = cpu_v_proj_seq + (seq_len - 1) * kv_dim; + const float *hybrid_v_last = hybrid_v_proj_seq + (seq_len - 1) * kv_dim; + for (int i = 0; i < trace_channel_count; i++) { + int channel = trace_channels[i]; + int head = channel / head_dim; + int offset = channel % head_dim; + int kv_head = head / q_per_kv; + int kv_channel = kv_head * head_dim + offset; + if (kv_channel < 0 || kv_channel >= kv_dim) continue; + + double self_diff = fabs((double)cpu_v_last[kv_channel] - (double)hybrid_v_last[kv_channel]); + int global_best_cpu_channel = -1; + double global_best_diff = INFINITY; + int local_best_shift = 0; + int local_best_cpu_channel = kv_channel; + double local_best_diff = INFINITY; + int radius = 8; + for (int cpu_channel = 0; cpu_channel < kv_dim; cpu_channel++) { + double diff = fabs((double)cpu_v_last[cpu_channel] - (double)hybrid_v_last[kv_channel]); + if (diff < global_best_diff) { + global_best_diff = diff; + global_best_cpu_channel = cpu_channel; + } + int shift = cpu_channel - kv_channel; + if (shift < -radius || shift > radius) continue; + if (diff < local_best_diff) { + local_best_diff = diff; + local_best_shift = shift; + local_best_cpu_channel = cpu_channel; + } + } + printf("trace_v_neighbor source=%s channel=%d head=%d offset=%d kv_head=%d kv_channel=%d self_diff=%.6f local_best_shift=%d local_best_cpu_channel=%d local_best_diff=%.6f global_best_cpu_channel=%d global_best_diff=%.6f\n", + source, channel, head, offset, kv_head, kv_channel, self_diff, + local_best_shift, local_best_cpu_channel, local_best_diff, + global_best_cpu_channel, global_best_diff); + } +} + +static inline float fp16_roundtrip(float x) { + return (float)(_Float16)x; +} + +static inline float sigmoid_scalar(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +static inline float silu_scalar(float x) { + return x * sigmoid_scalar(x); +} + +static inline float sigmoid_scalar_fp16_emulated(float x) { + float x16 = fp16_roundtrip(x); + float neg = fp16_roundtrip(-x16); + float expv = fp16_roundtrip(expf(neg)); + float denom = fp16_roundtrip(1.0f + expv); + return fp16_roundtrip(1.0f / denom); +} + +static inline float silu_scalar_fp16_emulated(float x) { + float x16 = fp16_roundtrip(x); + float sig = sigmoid_scalar_fp16_emulated(x16); + return fp16_roundtrip(x16 * sig); +} + +static inline float mul_scalar_fp16_emulated(float a, float b) { + float a16 = fp16_roundtrip(a); + float b16 = fp16_roundtrip(b); + return fp16_roundtrip(a16 * b16); +} + +static void apply_rope_text_inplace_local(float *q, + float *k, + int seq_len, + int n_q_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor) { + int rotary_dim = (int)(head_dim * partial_rotary_factor); + if (rotary_dim > head_dim) rotary_dim = head_dim; + if (rotary_dim % 2 != 0) rotary_dim -= 1; + if (rotary_dim <= 0) return; + + int half_rot = rotary_dim / 2; + float *inv_freq = (float *)malloc((size_t)half_rot * sizeof(float)); + if (!inv_freq) return; + + for (int i = 0; i < half_rot; i++) { + float exponent = (2.0f * (float)i) / (float)rotary_dim; + inv_freq[i] = 1.0f / powf(rope_theta, exponent); + } + + for (int pos = 0; pos < seq_len; pos++) { + for (int h = 0; h < n_q_head; h++) { + float *qh = q + pos * (n_q_head * head_dim) + h * head_dim; + for (int i = 0; i < half_rot; i++) { + float angle = (float)pos * inv_freq[i]; + float c = cosf(angle); + float s = sinf(angle); + float x0 = qh[i]; + float x1 = qh[i + half_rot]; + qh[i] = x0 * c - x1 * s; + qh[i + half_rot] = x1 * c + x0 * s; + } + } + for (int h = 0; h < n_kv_head; h++) { + float *kh = k + pos * (n_kv_head * head_dim) + h * head_dim; + for (int i = 0; i < half_rot; i++) { + float angle = (float)pos * inv_freq[i]; + float c = cosf(angle); + float s = sinf(angle); + float x0 = kh[i]; + float x1 = kh[i + half_rot]; + kh[i] = x0 * c - x1 * s; + kh[i + half_rot] = x1 * c + x0 * s; + } + } + } + + free(inv_freq); +} + +static int capture_attention_stages_from_projected_qkv(const float *q_proj_out_seq, + const float *k_proj_out_seq, + const float *v_proj_out_seq, + int seq_len, + const float *q_norm, + const float *k_norm, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + float *q_normed_out, + float *gate_sigmoid_out, + float *k_normed_out, + float *q_rope_out, + float *k_rope_out, + float *scores_out, + float *probs_out, + float *context_out, + float *gated_context_out) { + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int q_per_kv = n_head / n_kv_head; + const float scale = 1.0f / sqrtf((float)head_dim); + + float *gate_raw = (float *)calloc((size_t)seq_len * q_dim, sizeof(float)); + float *v_copy = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + float *scores = (float *)calloc((size_t)seq_len * seq_len, sizeof(float)); + float *probs = (float *)calloc((size_t)seq_len * seq_len, sizeof(float)); + float *qh = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + float *kh = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + float *vh = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + float *context_h = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + if (!gate_raw || !v_copy || !scores || !probs || !qh || !kh || !vh || !context_h) { + free(gate_raw); + free(v_copy); + free(scores); + free(probs); + free(qh); + free(kh); + free(vh); + free(context_h); + return 0; + } + + memcpy(v_copy, v_proj_out_seq, (size_t)seq_len * kv_dim * sizeof(float)); + memset(context_out, 0, (size_t)seq_len * q_dim * sizeof(float)); + memset(gated_context_out, 0, (size_t)seq_len * q_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + memcpy(q_normed_out + s * q_dim, + q_proj_out_seq + s * (q_dim * 2), + (size_t)q_dim * sizeof(float)); + memcpy(gate_raw + s * q_dim, + q_proj_out_seq + s * (q_dim * 2) + q_dim, + (size_t)q_dim * sizeof(float)); + memcpy(k_normed_out + s * kv_dim, + k_proj_out_seq + s * kv_dim, + (size_t)kv_dim * sizeof(float)); + } + + for (int s = 0; s < seq_len; s++) { + for (int h = 0; h < n_head; h++) { + orion_qwen_cpu_rmsnorm(q_normed_out + s * q_dim + h * head_dim, + q_norm, head_dim, 1e-6f, + q_normed_out + s * q_dim + h * head_dim); + } + for (int h = 0; h < n_kv_head; h++) { + orion_qwen_cpu_rmsnorm(k_normed_out + s * kv_dim + h * head_dim, + k_norm, head_dim, 1e-6f, + k_normed_out + s * kv_dim + h * head_dim); + } + for (int i = 0; i < q_dim; i++) { + gate_sigmoid_out[s * q_dim + i] = sigmoid_scalar(gate_raw[s * q_dim + i]); + } + } + + memcpy(q_rope_out, q_normed_out, (size_t)seq_len * q_dim * sizeof(float)); + memcpy(k_rope_out, k_normed_out, (size_t)seq_len * kv_dim * sizeof(float)); + apply_rope_text_inplace_local(q_rope_out, k_rope_out, seq_len, n_head, n_kv_head, head_dim, rope_theta, partial_rotary_factor); + + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + memset(scores, 0, (size_t)seq_len * seq_len * sizeof(float)); + memset(probs, 0, (size_t)seq_len * seq_len * sizeof(float)); + memset(context_h, 0, (size_t)seq_len * head_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + memcpy(qh + s * head_dim, q_rope_out + s * q_dim + h * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(kh + s * head_dim, k_rope_out + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(vh + s * head_dim, v_copy + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, seq_len, head_dim, + scale, qh, head_dim, kh, head_dim, + 0.0f, scores, seq_len); + + for (int i = 0; i < seq_len; i++) { + for (int j = i + 1; j < seq_len; j++) { + scores[i * seq_len + j] = -INFINITY; + } + float max_val = scores[i * seq_len]; + for (int j = 1; j <= i; j++) { + if (scores[i * seq_len + j] > max_val) max_val = scores[i * seq_len + j]; + } + float sum = 0.0f; + for (int j = 0; j <= i; j++) { + probs[i * seq_len + j] = expf(scores[i * seq_len + j] - max_val); + sum += probs[i * seq_len + j]; + } + for (int j = 0; j <= i; j++) { + probs[i * seq_len + j] /= sum; + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + seq_len, head_dim, seq_len, + 1.0f, probs, seq_len, vh, head_dim, + 0.0f, context_h, head_dim); + + memcpy(scores_out + (size_t)h * seq_len * seq_len, scores, (size_t)seq_len * seq_len * sizeof(float)); + memcpy(probs_out + (size_t)h * seq_len * seq_len, probs, (size_t)seq_len * seq_len * sizeof(float)); + for (int s = 0; s < seq_len; s++) { + memcpy(context_out + s * q_dim + h * head_dim, + context_h + s * head_dim, + (size_t)head_dim * sizeof(float)); + } + } + + for (int i = 0; i < seq_len * q_dim; i++) { + gated_context_out[i] = context_out[i] * gate_sigmoid_out[i]; + } + + free(gate_raw); + free(v_copy); + free(scores); + free(probs); + free(qh); + free(kh); + free(vh); + free(context_h); + return 1; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + OrionGPT2Tokenizer* tok = NULL; + OrionQwen35Manifest* manifest = NULL; + float *hidden = NULL; + float *hidden_next = NULL; + float *normed = NULL; + float *normed_fp16input = NULL; + float *cpu_q = NULL; + float *cpu_k = NULL; + float *cpu_v = NULL; + float *cpu_v_rms_only = NULL; + float *cpu_v_linear_fp16emu = NULL; + float *cpu_v_linear_fp16acc32 = NULL; + float *cpu_v_input_fp16 = NULL; + float *cpu_input_rms_ms = NULL; + float *cpu_input_rms_mse = NULL; + float *cpu_input_rms_rrms = NULL; + float *cpu_input_rms_xr = NULL; + float *ane_q = NULL; + float *ane_k = NULL; + float *ane_v = NULL; + float *ane_v_single = NULL; + float *ane_input_rms = NULL; + float *ane_input_rms_ms = NULL; + float *ane_input_rms_rrms = NULL; + float *ane_input_rms_rrms_pow_micro = NULL; + float *ane_input_rms_xr = NULL; + float *ane_v_linear_only = NULL; + float *ane_v_linear_only_fp16input = NULL; + float *cpu_attn = NULL; + float *ane_attn = NULL; + float *attn_gate_only = NULL; + float *attn_v_only = NULL; + float *attn_v_single_only = NULL; + float *attn_v_linear_only = NULL; + float *attn_v_linear_fp16emu = NULL; + float *attn_v_linear_fp16acc32 = NULL; + float *attn_v_input_fp16_cpu = NULL; + float *attn_v_input_fp16_ane = NULL; + float *attn_v_rms_only = NULL; + float *sigmoid_only_gated_context = NULL; + float *attn_sigmoid_only = NULL; + float *ane_q_token_major = NULL; + float *ane_k_token_major = NULL; + float *ane_v_token_major = NULL; + float *cpu_hidden_attn = NULL; + float *ane_hidden_attn = NULL; + float *gate_only_hidden_attn = NULL; + float *sigmoid_only_hidden_attn = NULL; + float *v_only_hidden_attn = NULL; + float *cpu_attn_q_normed = NULL; + float *ane_attn_q_normed = NULL; + float *cpu_attn_gate_sigmoid = NULL; + float *ane_attn_gate_sigmoid = NULL; + float *cpu_attn_k_normed = NULL; + float *ane_attn_k_normed = NULL; + float *cpu_attn_q_rope = NULL; + float *ane_attn_q_rope = NULL; + float *cpu_attn_k_rope = NULL; + float *ane_attn_k_rope = NULL; + float *cpu_attn_scores = NULL; + float *ane_attn_scores = NULL; + float *cpu_attn_probs = NULL; + float *ane_attn_probs = NULL; + float *cpu_attn_context = NULL; + float *ane_attn_context = NULL; + float *cpu_attn_gated_context = NULL; + float *ane_attn_gated_context = NULL; + float *cpu_ffn_rms = NULL; + float *ane_ffn_rms = NULL; + float *cpu_ffn_gate = NULL; + float *ane_ffn_gate = NULL; + float *cpu_ffn_up = NULL; + float *ane_ffn_up = NULL; + float *ane_ffn_up_only = NULL; + float *cpu_ffn_sigmoid = NULL; + float *cpu_ffn_sigmoid_fp16emu = NULL; + float *ane_ffn_sigmoid_micro = NULL; + float *ane_ffn_sigmoid_tanh_micro = NULL; + float *cpu_ffn_silu = NULL; + float *cpu_ffn_silu_fp16emu = NULL; + float *ane_ffn_silu = NULL; + float *ane_ffn_silu_only = NULL; + float *ane_ffn_silu_micro = NULL; + float *ane_ffn_silu_tanh_micro = NULL; + float *cpu_ffn_hidden = NULL; + float *cpu_ffn_hidden_fp16emu = NULL; + float *ane_ffn_hidden = NULL; + float *ane_ffn_hidden_only = NULL; + float *ane_ffn_mul_micro = NULL; + float *cpu_ffn_down = NULL; + float *ane_ffn_down_same_input = NULL; + float *cpu_ffn_final = NULL; + float *ane_ffn_final_same_input = NULL; + float *hybrid_layer_final = NULL; + float *input_ln = NULL; + float *post_ln = NULL; + float *q_proj = NULL; + float *k_proj = NULL; + float *v_proj = NULL; + float *o_proj = NULL; + float *q_norm = NULL; + float *k_norm = NULL; + float *gate_proj = NULL; + float *up_proj = NULL; + float *down_proj = NULL; + float *final_norm = NULL; + float *cpu_stage_last = NULL; + float *hybrid_stage_last = NULL; + OrionQwen35AneBridge bridge; + memset(&bridge, 0, sizeof(bridge)); + IOSurfaceRef ioIn = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + IOSurfaceRef ioVSingle = NULL; + IOSurfaceRef ioNormedIn = NULL; + IOSurfaceRef ioNormedFp16In = NULL; + IOSurfaceRef ioInputRmsOnly = NULL; + IOSurfaceRef ioInputRmsMs = NULL; + IOSurfaceRef ioInputRmsRrms = NULL; + IOSurfaceRef ioInputRmsMseIn = NULL; + IOSurfaceRef ioInputRmsRrmsPowMicro = NULL; + IOSurfaceRef ioInputRmsXr = NULL; + IOSurfaceRef ioVLinearOnly = NULL; + IOSurfaceRef ioVLinearOnlyFp16In = NULL; + IOSurfaceRef ioFfnIn = NULL; + IOSurfaceRef ioFfnRms = NULL; + IOSurfaceRef ioFfnGate = NULL; + IOSurfaceRef ioFfnUp = NULL; + IOSurfaceRef ioFfnUpOnly = NULL; + IOSurfaceRef ioFfnSigmoidMicroIn = NULL; + IOSurfaceRef ioFfnSigmoidMicroOut = NULL; + IOSurfaceRef ioFfnSigmoidTanhMicroOut = NULL; + IOSurfaceRef ioFfnSilu = NULL; + IOSurfaceRef ioFfnSiluOnly = NULL; + IOSurfaceRef ioFfnSiluMicroIn = NULL; + IOSurfaceRef ioFfnSiluMicroOut = NULL; + IOSurfaceRef ioFfnSiluTanhMicroOut = NULL; + IOSurfaceRef ioFfnHidden = NULL; + IOSurfaceRef ioFfnHiddenOnly = NULL; + IOSurfaceRef ioFfnMulMicroA = NULL; + IOSurfaceRef ioFfnMulMicroB = NULL; + IOSurfaceRef ioFfnMulMicroOut = NULL; + IOSurfaceRef ioFfnDownIn = NULL; + IOSurfaceRef ioFfnDown = NULL; + IOSurfaceRef ioHidden = NULL; + IOSurfaceRef ioFfnInMixed = NULL; + IOSurfaceRef ioHiddenMixed = NULL; + + if (argc < 3) { + fprintf(stderr, "usage: %s [prompt] [layer] [candidate_a] [candidate_b] [trace_dims] [trace_channels]\n", argv[0]); + return 2; + } + if (!orion_ane_init()) { + fprintf(stderr, "FAIL: orion_ane_init failed\n"); + return 3; + } + + const char *blob_dir = argv[1]; + NSString* tokDir = [NSString stringWithUTF8String:argv[2]]; + NSDictionary* meta = load_json([tokDir stringByAppendingPathComponent:@"meta.json"]); + if (!meta) { + fprintf(stderr, "FAIL: missing tokenizer meta.json\n"); + return 1; + } + NSString* regex = meta[@"regex_pattern"]; + NSString* vocabPath = [tokDir stringByAppendingPathComponent:@"vocab.json"]; + NSString* mergesPath = [tokDir stringByAppendingPathComponent:@"merges.txt"]; + tok = orion_gpt2_tokenizer_load_with_regex(vocabPath.UTF8String, mergesPath.UTF8String, regex.UTF8String); + if (!tok) { + fprintf(stderr, "FAIL: tokenizer load failed\n"); + return 1; + } + + manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const char *prompt = (argc >= 4) ? argv[3] : "사진"; + int layer = (argc >= 5) ? atoi(argv[4]) : 3; + int candidate_a = (argc >= 6) ? atoi(argv[5]) : -1; + int candidate_b = (argc >= 7) ? atoi(argv[6]) : -1; + int pair_enabled = (candidate_a >= 0 && candidate_b >= 0); + if (layer < 0 || layer >= manifest->n_layer) { + fprintf(stderr, "FAIL: invalid layer %d\n", layer); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + int token_ids[128] = {0}; + int seq_len = orion_gpt2_encode(tok, prompt, token_ids, 128); + if (seq_len <= 0) { + fprintf(stderr, "FAIL: prompt encode failed\n"); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int q_per_kv = n_head / n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int bucket = 32; + const int total = seq_len * d_model; + const int total_score = n_head * seq_len * seq_len; + const int total_q = seq_len * q_dim; + const int total_kv = seq_len * kv_dim; + const int total_ff = seq_len * d_ff; + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + const char *trace_dims_csv = (argc >= 8) ? argv[7] : ""; + const char *trace_channels_csv = (argc >= 9) ? argv[8] : ""; + int trace_dims[ORION_TRACE_DIM_MAX] = {0}; + int trace_dim_count = parse_dim_list(trace_dims_csv, trace_dims, ORION_TRACE_DIM_MAX, d_model); + int trace_channels[ORION_TRACE_DIM_MAX] = {0}; + int trace_channel_count = parse_dim_list(trace_channels_csv, trace_channels, ORION_TRACE_DIM_MAX, q_dim); + const int q_proj_uses_cpu = use_cpu_q_proj_override(); + const int q_query_uses_cpu = use_cpu_q_query_override(); + const int q_gate_uses_cpu = use_cpu_q_gate_override(); + const int v_proj_uses_cpu = use_cpu_v_proj_override(); + int v_proj_cpu_channels[ORION_TRACE_DIM_MAX] = {0}; + int v_proj_cpu_channel_count = 0; + if (!v_proj_uses_cpu) { + v_proj_cpu_channel_count = load_v_proj_cpu_channel_override(v_proj_cpu_channels, ORION_TRACE_DIM_MAX, q_dim); + } + NSString *v_linear_only_tmp_dir = nil; + NSString *v_linear_only_source_blob_path = nil; + NSString *v_linear_only_data_path = nil; + NSString *v_linear_only_weight_path = nil; + NSData *v_linear_only_source_blob = nil; + NSData *v_linear_only_source_payload = nil; + NSData *v_linear_only_data_blob = nil; + NSData *v_linear_only_weight_blob = nil; + OrionByteCompareResult v_linear_only_data_vs_source = {0}; + OrionByteCompareResult v_linear_only_weight_vs_source = {0}; + OrionByteCompareResult v_linear_only_data_vs_weight = {0}; + OrionByteCompareResult v_linear_only_data_vs_source_payload = {0}; + + hidden = (float *)calloc((size_t)total, sizeof(float)); + hidden_next = (float *)calloc((size_t)total, sizeof(float)); + normed = (float *)calloc((size_t)total, sizeof(float)); + normed_fp16input = (float *)calloc((size_t)total, sizeof(float)); + cpu_q = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + cpu_k = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v_rms_only = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v_linear_fp16emu = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v_linear_fp16acc32 = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v_input_fp16 = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_input_rms_ms = (float *)calloc((size_t)seq_len, sizeof(float)); + cpu_input_rms_mse = (float *)calloc((size_t)seq_len, sizeof(float)); + cpu_input_rms_rrms = (float *)calloc((size_t)seq_len, sizeof(float)); + cpu_input_rms_xr = (float *)calloc((size_t)total, sizeof(float)); + ane_q = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + ane_k = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_v = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_v_single = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_input_rms = (float *)calloc((size_t)total, sizeof(float)); + ane_input_rms_ms = (float *)calloc((size_t)seq_len, sizeof(float)); + ane_input_rms_rrms = (float *)calloc((size_t)seq_len, sizeof(float)); + ane_input_rms_rrms_pow_micro = (float *)calloc((size_t)seq_len, sizeof(float)); + ane_input_rms_xr = (float *)calloc((size_t)total, sizeof(float)); + ane_v_linear_only = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_v_linear_only_fp16input = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_q_token_major = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + ane_k_token_major = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_v_token_major = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_attn = (float *)calloc((size_t)total, sizeof(float)); + ane_attn = (float *)calloc((size_t)total, sizeof(float)); + attn_gate_only = (float *)calloc((size_t)total, sizeof(float)); + attn_v_only = (float *)calloc((size_t)total, sizeof(float)); + attn_v_single_only = (float *)calloc((size_t)total, sizeof(float)); + attn_v_linear_only = (float *)calloc((size_t)total, sizeof(float)); + attn_v_linear_fp16emu = (float *)calloc((size_t)total, sizeof(float)); + attn_v_linear_fp16acc32 = (float *)calloc((size_t)total, sizeof(float)); + attn_v_input_fp16_cpu = (float *)calloc((size_t)total, sizeof(float)); + attn_v_input_fp16_ane = (float *)calloc((size_t)total, sizeof(float)); + attn_v_rms_only = (float *)calloc((size_t)total, sizeof(float)); + sigmoid_only_gated_context = (float *)calloc((size_t)total_q, sizeof(float)); + attn_sigmoid_only = (float *)calloc((size_t)total, sizeof(float)); + cpu_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + ane_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + gate_only_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + sigmoid_only_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + v_only_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + cpu_attn_q_normed = (float *)calloc((size_t)total_q, sizeof(float)); + ane_attn_q_normed = (float *)calloc((size_t)total_q, sizeof(float)); + cpu_attn_gate_sigmoid = (float *)calloc((size_t)total_q, sizeof(float)); + ane_attn_gate_sigmoid = (float *)calloc((size_t)total_q, sizeof(float)); + cpu_attn_k_normed = (float *)calloc((size_t)total_kv, sizeof(float)); + ane_attn_k_normed = (float *)calloc((size_t)total_kv, sizeof(float)); + cpu_attn_q_rope = (float *)calloc((size_t)total_q, sizeof(float)); + ane_attn_q_rope = (float *)calloc((size_t)total_q, sizeof(float)); + cpu_attn_k_rope = (float *)calloc((size_t)total_kv, sizeof(float)); + ane_attn_k_rope = (float *)calloc((size_t)total_kv, sizeof(float)); + cpu_attn_scores = (float *)calloc((size_t)total_score, sizeof(float)); + ane_attn_scores = (float *)calloc((size_t)total_score, sizeof(float)); + cpu_attn_probs = (float *)calloc((size_t)total_score, sizeof(float)); + ane_attn_probs = (float *)calloc((size_t)total_score, sizeof(float)); + cpu_attn_context = (float *)calloc((size_t)total_q, sizeof(float)); + ane_attn_context = (float *)calloc((size_t)total_q, sizeof(float)); + cpu_attn_gated_context = (float *)calloc((size_t)total_q, sizeof(float)); + ane_attn_gated_context = (float *)calloc((size_t)total_q, sizeof(float)); + cpu_ffn_rms = (float *)calloc((size_t)total, sizeof(float)); + ane_ffn_rms = (float *)calloc((size_t)total, sizeof(float)); + cpu_ffn_gate = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_gate = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_up = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_up = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_up_only = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_sigmoid = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_sigmoid_fp16emu = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_sigmoid_micro = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_sigmoid_tanh_micro = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_silu = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_silu_fp16emu = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_silu = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_silu_only = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_silu_micro = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_silu_tanh_micro = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_hidden = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_hidden_fp16emu = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_hidden = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_hidden_only = (float *)calloc((size_t)total_ff, sizeof(float)); + ane_ffn_mul_micro = (float *)calloc((size_t)total_ff, sizeof(float)); + cpu_ffn_down = (float *)calloc((size_t)total, sizeof(float)); + ane_ffn_down_same_input = (float *)calloc((size_t)total, sizeof(float)); + cpu_ffn_final = (float *)calloc((size_t)total, sizeof(float)); + ane_ffn_final_same_input = (float *)calloc((size_t)total, sizeof(float)); + hybrid_layer_final = (float *)calloc((size_t)total, sizeof(float)); + cpu_stage_last = (float *)calloc((size_t)d_model, sizeof(float)); + hybrid_stage_last = (float *)calloc((size_t)d_model, sizeof(float)); + if (!hidden || !hidden_next || !normed || !normed_fp16input || !cpu_q || !cpu_k || !cpu_v || !cpu_v_rms_only || !cpu_v_linear_fp16emu || !cpu_v_linear_fp16acc32 || !cpu_v_input_fp16 || + !cpu_input_rms_ms || !cpu_input_rms_mse || !cpu_input_rms_rrms || !cpu_input_rms_xr || + !ane_q || !ane_k || !ane_v || !ane_v_single || !ane_input_rms || + !ane_input_rms_ms || !ane_input_rms_rrms || !ane_input_rms_rrms_pow_micro || !ane_input_rms_xr || + !ane_v_linear_only || !ane_v_linear_only_fp16input || + !ane_q_token_major || !ane_k_token_major || !ane_v_token_major || + !cpu_attn || !ane_attn || !attn_gate_only || !attn_v_only || !attn_v_single_only || !attn_v_linear_only || !attn_v_linear_fp16emu || !attn_v_linear_fp16acc32 || !attn_v_input_fp16_cpu || !attn_v_input_fp16_ane || !attn_v_rms_only || + !sigmoid_only_gated_context || !attn_sigmoid_only || + !cpu_hidden_attn || !ane_hidden_attn || !gate_only_hidden_attn || !sigmoid_only_hidden_attn || !v_only_hidden_attn || + !cpu_attn_q_normed || !ane_attn_q_normed || !cpu_attn_gate_sigmoid || !ane_attn_gate_sigmoid || + !cpu_attn_k_normed || !ane_attn_k_normed || !cpu_attn_q_rope || !ane_attn_q_rope || + !cpu_attn_k_rope || !ane_attn_k_rope || !cpu_attn_scores || !ane_attn_scores || + !cpu_attn_probs || !ane_attn_probs || !cpu_attn_context || !ane_attn_context || + !cpu_attn_gated_context || !ane_attn_gated_context || + !cpu_ffn_rms || !ane_ffn_rms || !cpu_ffn_gate || !ane_ffn_gate || + !cpu_ffn_up || !ane_ffn_up || !ane_ffn_up_only || + !cpu_ffn_sigmoid || !cpu_ffn_sigmoid_fp16emu || !ane_ffn_sigmoid_micro || !ane_ffn_sigmoid_tanh_micro || + !cpu_ffn_silu || !cpu_ffn_silu_fp16emu || !ane_ffn_silu || !ane_ffn_silu_only || !ane_ffn_silu_micro || !ane_ffn_silu_tanh_micro || + !cpu_ffn_hidden || !cpu_ffn_hidden_fp16emu || !ane_ffn_hidden || !ane_ffn_hidden_only || !ane_ffn_mul_micro || + !cpu_ffn_down || !ane_ffn_down_same_input || + !cpu_ffn_final || !ane_ffn_final_same_input || !hybrid_layer_final || + !cpu_stage_last || !hybrid_stage_last) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + if (!load_embeddings(blob_dir, manifest, token_ids, seq_len, hidden)) { + fprintf(stderr, "FAIL: embedding load failed\n"); + goto fail; + } + for (int i = 0; i < layer; i++) { + if (!apply_cpu_layer(blob_dir, manifest, i, hidden, seq_len, hidden_next)) { + fprintf(stderr, "FAIL: pre-layer cpu apply failed at layer %d\n", i); + goto fail; + } + float *tmp = hidden; + hidden = hidden_next; + hidden_next = tmp; + } + + input_ln = load_exact(blob_dir, layer, "input_layernorm.bin", d_model); + post_ln = load_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + q_proj = load_exact(blob_dir, layer, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + k_proj = load_exact(blob_dir, layer, "self_attn_k_proj.bin", kv_dim * d_model); + v_proj = load_exact(blob_dir, layer, "self_attn_v_proj.bin", kv_dim * d_model); + o_proj = load_exact(blob_dir, layer, "self_attn_o_proj.bin", d_model * q_dim); + q_norm = load_exact(blob_dir, layer, "self_attn_q_norm.bin", head_dim); + k_norm = load_exact(blob_dir, layer, "self_attn_k_norm.bin", head_dim); + gate_proj = load_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + up_proj = load_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + down_proj = load_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!input_ln || !post_ln || !q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm || + !gate_proj || !up_proj || !down_proj || !final_norm) { + fprintf(stderr, "FAIL: missing layer weights\n"); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + double sumsq = 0.0; + for (int i = 0; i < d_model; i++) { + double x = (double)hidden[s * d_model + i]; + sumsq += x * x; + } + cpu_input_rms_ms[s] = (float)(sumsq / (double)d_model); + cpu_input_rms_mse[s] = cpu_input_rms_ms[s] + 1e-6f; + cpu_input_rms_rrms[s] = powf(cpu_input_rms_mse[s], -0.5f); + for (int i = 0; i < d_model; i++) { + cpu_input_rms_xr[s * d_model + i] = hidden[s * d_model + i] * cpu_input_rms_rrms[s]; + } + } + for (int i = 0; i < total; i++) { + normed_fp16input[i] = fp16_roundtrip(normed[i]); + } + + cpu_linear_batch(normed, seq_len, q_proj, d_model, q_dim * 2, cpu_q); + cpu_linear_batch(normed, seq_len, k_proj, d_model, kv_dim, cpu_k); + cpu_linear_batch(normed, seq_len, v_proj, d_model, kv_dim, cpu_v); + cpu_linear_batch_fp16_emulated(normed, seq_len, v_proj, d_model, kv_dim, cpu_v_linear_fp16emu); + cpu_linear_batch_fp16in_fp32acc_fp16out(normed, seq_len, v_proj, d_model, kv_dim, cpu_v_linear_fp16acc32); + cpu_linear_batch(normed_fp16input, seq_len, v_proj, d_model, kv_dim, cpu_v_input_fp16); + + if (!bridge_init(&bridge, [NSString stringWithUTF8String:blob_dir], layer, bucket, manifest)) { + fprintf(stderr, "FAIL: bridge_init failed\n"); + goto fail; + } + + v_linear_only_tmp_dir = orion_program_tmp_dir(bridge.prog_v_linear_only); + v_linear_only_source_blob_path = [NSString stringWithFormat:@"%s/layer%d/self_attn_v_proj.bin", blob_dir, layer]; + if (v_linear_only_tmp_dir) { + v_linear_only_data_path = [v_linear_only_tmp_dir stringByAppendingPathComponent:@"data"]; + v_linear_only_weight_path = [v_linear_only_tmp_dir stringByAppendingPathComponent: + [NSString stringWithFormat:@"layer%d/self_attn_v_proj.bin", layer]]; + } + v_linear_only_source_blob = [NSData dataWithContentsOfFile:v_linear_only_source_blob_path]; + v_linear_only_source_payload = load_blob_payload_data(v_linear_only_source_blob_path); + v_linear_only_data_blob = v_linear_only_data_path ? [NSData dataWithContentsOfFile:v_linear_only_data_path] : nil; + v_linear_only_weight_blob = v_linear_only_weight_path ? [NSData dataWithContentsOfFile:v_linear_only_weight_path] : nil; + v_linear_only_data_vs_source = compare_nsdata(v_linear_only_data_blob, v_linear_only_source_blob); + v_linear_only_weight_vs_source = compare_nsdata(v_linear_only_weight_blob, v_linear_only_source_blob); + v_linear_only_data_vs_weight = compare_nsdata(v_linear_only_data_blob, v_linear_only_weight_blob); + v_linear_only_data_vs_source_payload = compare_nsdata(v_linear_only_data_blob, v_linear_only_source_payload); + + ioIn = make_cpu_seq_input_surface(hidden, seq_len, bucket, d_model); + ioQ = make_f32_surface((q_dim * 2) * bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bucket, 0.0f); + ioVSingle = make_f32_surface(kv_dim * bucket, 0.0f); + ioNormedIn = make_cpu_seq_input_surface(normed, seq_len, bucket, d_model); + ioNormedFp16In = make_cpu_seq_input_surface(normed_fp16input, seq_len, bucket, d_model); + ioInputRmsOnly = make_f32_surface(d_model * bucket, 0.0f); + ioInputRmsMs = make_f32_surface(bucket, 0.0f); + ioInputRmsRrms = make_f32_surface(bucket, 0.0f); + ioInputRmsMseIn = make_cpu_seq_input_surface(cpu_input_rms_mse, seq_len, bucket, 1); + ioInputRmsRrmsPowMicro = make_f32_surface(bucket, 0.0f); + ioInputRmsXr = make_f32_surface(d_model * bucket, 0.0f); + ioVLinearOnly = make_f32_surface(kv_dim * bucket, 0.0f); + ioVLinearOnlyFp16In = make_f32_surface(kv_dim * bucket, 0.0f); + IOSurfaceRef insQ[] = {ioIn}; + IOSurfaceRef insNormed[] = {ioNormedIn}; + IOSurfaceRef insNormedFp16[] = {ioNormedFp16In}; + IOSurfaceRef outsQ[] = {ioQ}; + IOSurfaceRef outsKV[] = {ioK, ioV}; + IOSurfaceRef outsVSingle[] = {ioVSingle}; + IOSurfaceRef outsInputRmsOnly[] = {ioInputRmsOnly}; + IOSurfaceRef outsInputRmsMs[] = {ioInputRmsMs}; + IOSurfaceRef outsInputRmsRrms[] = {ioInputRmsRrms}; + IOSurfaceRef outsInputRmsXr[] = {ioInputRmsXr}; + IOSurfaceRef insInputRmsMse[] = {ioInputRmsMseIn}; + IOSurfaceRef outsInputRmsRrmsPowMicro[] = {ioInputRmsRrmsPowMicro}; + IOSurfaceRef outsVLinearOnly[] = {ioVLinearOnly}; + IOSurfaceRef outsVLinearOnlyFp16In[] = {ioVLinearOnlyFp16In}; + if (!orion_eval(bridge.prog_q, insQ, 1, outsQ, 1) || !orion_eval(bridge.prog_kv, insQ, 1, outsKV, 2)) { + fprintf(stderr, "FAIL: ANE q/kv eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_v_only, insQ, 1, outsVSingle, 1)) { + fprintf(stderr, "FAIL: ANE v-only eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_input_rms_only, insQ, 1, outsInputRmsOnly, 1)) { + fprintf(stderr, "FAIL: ANE input-rms-only eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_input_rms_ms, insQ, 1, outsInputRmsMs, 1)) { + fprintf(stderr, "FAIL: ANE input-rms-ms eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_input_rms_rrms, insQ, 1, outsInputRmsRrms, 1)) { + fprintf(stderr, "FAIL: ANE input-rms-rrms eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_input_rms_xr, insQ, 1, outsInputRmsXr, 1)) { + fprintf(stderr, "FAIL: ANE input-rms-xr eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_rrms_pow_micro, insInputRmsMse, 1, outsInputRmsRrmsPowMicro, 1)) { + fprintf(stderr, "FAIL: ANE rrms-pow-micro eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_v_linear_only, insNormed, 1, outsVLinearOnly, 1)) { + fprintf(stderr, "FAIL: ANE v-linear-only eval failed\n"); + goto fail; + } + if (!orion_eval(bridge.prog_v_linear_only, insNormedFp16, 1, outsVLinearOnlyFp16In, 1)) { + fprintf(stderr, "FAIL: ANE v-linear-only eval failed on fp16-rounded input\n"); + goto fail; + } + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bucket, ane_q); + read_ane_surface_prefix(ioK, kv_dim, seq_len, bucket, ane_k); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bucket, ane_v); + read_ane_surface_prefix(ioVSingle, kv_dim, seq_len, bucket, ane_v_single); + read_ane_surface_prefix(ioInputRmsOnly, d_model, seq_len, bucket, ane_input_rms); + read_ane_surface_prefix(ioInputRmsMs, 1, seq_len, bucket, ane_input_rms_ms); + read_ane_surface_prefix(ioInputRmsRrms, 1, seq_len, bucket, ane_input_rms_rrms); + read_ane_surface_prefix(ioInputRmsRrmsPowMicro, 1, seq_len, bucket, ane_input_rms_rrms_pow_micro); + read_ane_surface_prefix(ioInputRmsXr, d_model, seq_len, bucket, ane_input_rms_xr); + read_ane_surface_prefix(ioVLinearOnly, kv_dim, seq_len, bucket, ane_v_linear_only); + read_ane_surface_prefix(ioVLinearOnlyFp16In, kv_dim, seq_len, bucket, ane_v_linear_only_fp16input); + read_ane_surface_prefix_token_major(ioQ, q_dim * 2, seq_len, ane_q_token_major); + read_ane_surface_prefix_token_major(ioK, kv_dim, seq_len, ane_k_token_major); + read_ane_surface_prefix_token_major(ioV, kv_dim, seq_len, ane_v_token_major); + if (q_proj_uses_cpu || q_query_uses_cpu || q_gate_uses_cpu) { + for (int s = 0; s < seq_len; s++) { + float *dst = ane_q + (size_t)s * (q_dim * 2); + float *dst_token_major = ane_q_token_major + (size_t)s * (q_dim * 2); + const float *src = cpu_q + (size_t)s * (q_dim * 2); + if (q_proj_uses_cpu || q_query_uses_cpu) { + memcpy(dst, src, (size_t)q_dim * sizeof(float)); + memcpy(dst_token_major, src, (size_t)q_dim * sizeof(float)); + } + if (q_proj_uses_cpu || q_gate_uses_cpu) { + memcpy(dst + q_dim, src + q_dim, (size_t)q_dim * sizeof(float)); + memcpy(dst_token_major + q_dim, src + q_dim, (size_t)q_dim * sizeof(float)); + } + } + } + if (v_proj_uses_cpu) { + memcpy(ane_v, cpu_v, (size_t)seq_len * kv_dim * sizeof(float)); + memcpy(ane_v_single, cpu_v, (size_t)seq_len * kv_dim * sizeof(float)); + memcpy(ane_v_linear_only, cpu_v, (size_t)seq_len * kv_dim * sizeof(float)); + memcpy(ane_v_linear_only_fp16input, cpu_v, (size_t)seq_len * kv_dim * sizeof(float)); + memcpy(ane_v_token_major, cpu_v, (size_t)seq_len * kv_dim * sizeof(float)); + } else if (v_proj_cpu_channel_count > 0) { + apply_v_proj_cpu_channel_overrides(ane_v, cpu_v, seq_len, kv_dim, head_dim, n_head / n_kv_head, + v_proj_cpu_channels, v_proj_cpu_channel_count); + apply_v_proj_cpu_channel_overrides(ane_v_single, cpu_v, seq_len, kv_dim, head_dim, n_head / n_kv_head, + v_proj_cpu_channels, v_proj_cpu_channel_count); + apply_v_proj_cpu_channel_overrides(ane_v_linear_only, cpu_v, seq_len, kv_dim, head_dim, n_head / n_kv_head, + v_proj_cpu_channels, v_proj_cpu_channel_count); + apply_v_proj_cpu_channel_overrides(ane_v_linear_only_fp16input, cpu_v, seq_len, kv_dim, head_dim, n_head / n_kv_head, + v_proj_cpu_channels, v_proj_cpu_channel_count); + apply_v_proj_cpu_channel_overrides(ane_v_token_major, cpu_v, seq_len, kv_dim, head_dim, n_head / n_kv_head, + v_proj_cpu_channels, v_proj_cpu_channel_count); + } + cpu_linear_batch(ane_input_rms, seq_len, v_proj, d_model, kv_dim, cpu_v_rms_only); + + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + cpu_attn + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + ane_q, ane_k, ane_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + ane_attn + ); + float *cpu_q_with_ane_gate = (float *)malloc((size_t)seq_len * (q_dim * 2) * sizeof(float)); + if (!cpu_q_with_ane_gate) { + fprintf(stderr, "FAIL: allocation failed for cpu_q_with_ane_gate\n"); + goto fail; + } + memcpy(cpu_q_with_ane_gate, cpu_q, (size_t)seq_len * (q_dim * 2) * sizeof(float)); + for (int s = 0; s < seq_len; s++) { + memcpy(cpu_q_with_ane_gate + s * (q_dim * 2) + q_dim, + ane_q + s * (q_dim * 2) + q_dim, + (size_t)q_dim * sizeof(float)); + } + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q_with_ane_gate, cpu_k, cpu_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_gate_only + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, ane_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_only + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, ane_v_single, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_single_only + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, ane_v_linear_only, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_linear_only + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v_linear_fp16emu, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_linear_fp16emu + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v_linear_fp16acc32, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_linear_fp16acc32 + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v_input_fp16, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_input_fp16_cpu + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, ane_v_linear_only_fp16input, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_input_fp16_ane + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v_rms_only, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn_v_rms_only + ); + free(cpu_q_with_ane_gate); + if (!capture_attention_stages_from_projected_qkv(cpu_q, cpu_k, cpu_v, seq_len, + q_norm, k_norm, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + cpu_attn_q_normed, cpu_attn_gate_sigmoid, cpu_attn_k_normed, + cpu_attn_q_rope, cpu_attn_k_rope, cpu_attn_scores, cpu_attn_probs, + cpu_attn_context, cpu_attn_gated_context)) { + fprintf(stderr, "FAIL: cpu attention stage capture failed\n"); + goto fail; + } + if (!capture_attention_stages_from_projected_qkv(ane_q, ane_k, ane_v, seq_len, + q_norm, k_norm, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + ane_attn_q_normed, ane_attn_gate_sigmoid, ane_attn_k_normed, + ane_attn_q_rope, ane_attn_k_rope, ane_attn_scores, ane_attn_probs, + ane_attn_context, ane_attn_gated_context)) { + fprintf(stderr, "FAIL: ane attention stage capture failed\n"); + goto fail; + } + for (int i = 0; i < total_q; i++) { + sigmoid_only_gated_context[i] = cpu_attn_context[i] * ane_attn_gate_sigmoid[i]; + } + cpu_linear_batch(sigmoid_only_gated_context, seq_len, o_proj, q_dim, d_model, attn_sigmoid_only); + + for (int i = 0; i < total; i++) { + cpu_hidden_attn[i] = hidden[i] + cpu_attn[i]; + ane_hidden_attn[i] = hidden[i] + ane_attn[i]; + gate_only_hidden_attn[i] = hidden[i] + attn_gate_only[i]; + sigmoid_only_hidden_attn[i] = hidden[i] + attn_sigmoid_only[i]; + v_only_hidden_attn[i] = hidden[i] + attn_v_only[i]; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + s * d_model, post_ln, d_model, 1e-6f, cpu_ffn_rms + s * d_model); + } + cpu_linear_batch(cpu_ffn_rms, seq_len, gate_proj, d_model, d_ff, cpu_ffn_gate); + cpu_linear_batch(cpu_ffn_rms, seq_len, up_proj, d_model, d_ff, cpu_ffn_up); + for (int i = 0; i < total_ff; i++) { + cpu_ffn_sigmoid[i] = sigmoid_scalar(cpu_ffn_gate[i]); + cpu_ffn_sigmoid_fp16emu[i] = sigmoid_scalar_fp16_emulated(cpu_ffn_gate[i]); + cpu_ffn_silu[i] = silu_scalar(cpu_ffn_gate[i]); + cpu_ffn_silu_fp16emu[i] = silu_scalar_fp16_emulated(cpu_ffn_gate[i]); + cpu_ffn_hidden[i] = cpu_ffn_silu[i] * cpu_ffn_up[i]; + cpu_ffn_hidden_fp16emu[i] = mul_scalar_fp16_emulated(cpu_ffn_silu_fp16emu[i], cpu_ffn_up[i]); + } + cpu_linear_batch(cpu_ffn_hidden, seq_len, down_proj, d_ff, d_model, cpu_ffn_down); + for (int i = 0; i < total; i++) cpu_ffn_final[i] = cpu_hidden_attn[i] + cpu_ffn_down[i]; + + ioFfnIn = make_cpu_seq_input_surface(cpu_hidden_attn, seq_len, bucket, d_model); + ioFfnRms = make_f32_surface(d_model * bucket, 0.0f); + ioFfnGate = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnUp = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnUpOnly = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnSilu = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnSiluOnly = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnHidden = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnHiddenOnly = make_f32_surface(d_ff * bucket, 0.0f); + ioHidden = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceRef insFFN[] = {ioFfnIn}; + IOSurfaceRef outsFFNRms[] = {ioFfnRms}; + IOSurfaceRef outsFFNGateUp[] = {ioFfnGate, ioFfnUp, ioFfnSilu, ioFfnHidden}; + IOSurfaceRef outsFFNUpOnly[] = {ioFfnUpOnly}; + IOSurfaceRef outsFFNSiluOnly[] = {ioFfnSiluOnly}; + IOSurfaceRef outsFFNHiddenOnly[] = {ioFfnHiddenOnly}; + IOSurfaceRef outsFFN[] = {ioHidden}; + if (!orion_eval(bridge.prog_ffn_rms, insFFN, 1, outsFFNRms, 1)) { + fprintf(stderr, "FAIL: ANE ffn rms eval failed on cpu_hidden_attn\n"); + goto fail; + } + if (!orion_eval(bridge.prog_ffn_gateup, insFFN, 1, outsFFNGateUp, 4)) { + fprintf(stderr, "FAIL: ANE ffn gate/up eval failed on cpu_hidden_attn\n"); + goto fail; + } + if (!orion_eval(bridge.prog_ffn_up_only, insFFN, 1, outsFFNUpOnly, 1)) { + fprintf(stderr, "FAIL: ANE ffn up-only eval failed on cpu_hidden_attn\n"); + goto fail; + } + if (!orion_eval(bridge.prog_ffn_silu_only, insFFN, 1, outsFFNSiluOnly, 1)) { + fprintf(stderr, "FAIL: ANE ffn silu-only eval failed on cpu_hidden_attn\n"); + goto fail; + } + if (!orion_eval(bridge.prog_ffn_hidden_only, insFFN, 1, outsFFNHiddenOnly, 1)) { + fprintf(stderr, "FAIL: ANE ffn hidden-only eval failed on cpu_hidden_attn\n"); + goto fail; + } + if (!orion_eval(bridge.prog_ffn, insFFN, 1, outsFFN, 1)) { + fprintf(stderr, "FAIL: ANE ffn eval failed on cpu_hidden_attn\n"); + goto fail; + } + read_ane_surface_prefix(ioFfnRms, d_model, seq_len, bucket, ane_ffn_rms); + read_ane_surface_prefix(ioFfnGate, d_ff, seq_len, bucket, ane_ffn_gate); + read_ane_surface_prefix(ioFfnUp, d_ff, seq_len, bucket, ane_ffn_up); + read_ane_surface_prefix(ioFfnUpOnly, d_ff, seq_len, bucket, ane_ffn_up_only); + read_ane_surface_prefix(ioFfnSilu, d_ff, seq_len, bucket, ane_ffn_silu); + read_ane_surface_prefix(ioFfnSiluOnly, d_ff, seq_len, bucket, ane_ffn_silu_only); + read_ane_surface_prefix(ioFfnHidden, d_ff, seq_len, bucket, ane_ffn_hidden); + read_ane_surface_prefix(ioFfnHiddenOnly, d_ff, seq_len, bucket, ane_ffn_hidden_only); + read_ane_surface_prefix(ioHidden, d_model, seq_len, bucket, ane_ffn_final_same_input); + + ioFfnSigmoidMicroIn = make_cpu_seq_input_surface(cpu_ffn_gate, seq_len, bucket, d_ff); + ioFfnSigmoidMicroOut = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnSigmoidTanhMicroOut = make_f32_surface(d_ff * bucket, 0.0f); + IOSurfaceRef insSigmoidMicro[] = {ioFfnSigmoidMicroIn}; + IOSurfaceRef outsSigmoidMicro[] = {ioFfnSigmoidMicroOut}; + IOSurfaceRef outsSigmoidTanhMicro[] = {ioFfnSigmoidTanhMicroOut}; + if (!orion_eval(bridge.prog_sigmoid_micro, insSigmoidMicro, 1, outsSigmoidMicro, 1)) { + fprintf(stderr, "FAIL: ANE sigmoid micro eval failed on cpu_ffn_gate\n"); + goto fail; + } + if (!orion_eval(bridge.prog_sigmoid_tanh_micro, insSigmoidMicro, 1, outsSigmoidTanhMicro, 1)) { + fprintf(stderr, "FAIL: ANE sigmoid tanh micro eval failed on cpu_ffn_gate\n"); + goto fail; + } + read_ane_surface_prefix(ioFfnSigmoidMicroOut, d_ff, seq_len, bucket, ane_ffn_sigmoid_micro); + read_ane_surface_prefix(ioFfnSigmoidTanhMicroOut, d_ff, seq_len, bucket, ane_ffn_sigmoid_tanh_micro); + + ioFfnSiluMicroIn = make_cpu_seq_input_surface(cpu_ffn_gate, seq_len, bucket, d_ff); + ioFfnSiluMicroOut = make_f32_surface(d_ff * bucket, 0.0f); + ioFfnSiluTanhMicroOut = make_f32_surface(d_ff * bucket, 0.0f); + IOSurfaceRef insSiluMicro[] = {ioFfnSiluMicroIn}; + IOSurfaceRef outsSiluMicro[] = {ioFfnSiluMicroOut}; + IOSurfaceRef outsSiluTanhMicro[] = {ioFfnSiluTanhMicroOut}; + if (!orion_eval(bridge.prog_silu_micro, insSiluMicro, 1, outsSiluMicro, 1)) { + fprintf(stderr, "FAIL: ANE silu micro eval failed on cpu_ffn_gate\n"); + goto fail; + } + if (!orion_eval(bridge.prog_silu_tanh_micro, insSiluMicro, 1, outsSiluTanhMicro, 1)) { + fprintf(stderr, "FAIL: ANE silu tanh micro eval failed on cpu_ffn_gate\n"); + goto fail; + } + read_ane_surface_prefix(ioFfnSiluMicroOut, d_ff, seq_len, bucket, ane_ffn_silu_micro); + read_ane_surface_prefix(ioFfnSiluTanhMicroOut, d_ff, seq_len, bucket, ane_ffn_silu_tanh_micro); + + ioFfnMulMicroA = make_cpu_seq_input_surface(cpu_ffn_silu, seq_len, bucket, d_ff); + ioFfnMulMicroB = make_cpu_seq_input_surface(cpu_ffn_up, seq_len, bucket, d_ff); + ioFfnMulMicroOut = make_f32_surface(d_ff * bucket, 0.0f); + IOSurfaceRef insMulMicro[] = {ioFfnMulMicroA, ioFfnMulMicroB}; + IOSurfaceRef outsMulMicro[] = {ioFfnMulMicroOut}; + if (!orion_eval(bridge.prog_mul_micro, insMulMicro, 2, outsMulMicro, 1)) { + fprintf(stderr, "FAIL: ANE mul micro eval failed on cpu_ffn_silu/cpu_ffn_up\n"); + goto fail; + } + read_ane_surface_prefix(ioFfnMulMicroOut, d_ff, seq_len, bucket, ane_ffn_mul_micro); + + ioFfnDownIn = make_cpu_seq_input_surface(cpu_ffn_hidden, seq_len, bucket, d_ff); + ioFfnDown = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceRef insFFNDown[] = {ioFfnDownIn}; + IOSurfaceRef outsFFNDown[] = {ioFfnDown}; + if (!orion_eval(bridge.prog_ffn_down, insFFNDown, 1, outsFFNDown, 1)) { + fprintf(stderr, "FAIL: ANE ffn down eval failed on cpu_ffn_hidden\n"); + goto fail; + } + read_ane_surface_prefix(ioFfnDown, d_model, seq_len, bucket, ane_ffn_down_same_input); + + ioFfnInMixed = make_cpu_seq_input_surface(ane_hidden_attn, seq_len, bucket, d_model); + ioHiddenMixed = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceRef insFFNMixed[] = {ioFfnInMixed}; + IOSurfaceRef outsFFNMixed[] = {ioHiddenMixed}; + if (!orion_eval(bridge.prog_ffn, insFFNMixed, 1, outsFFNMixed, 1)) { + fprintf(stderr, "FAIL: ANE ffn eval failed on ane_hidden_attn\n"); + goto fail; + } + read_ane_surface_prefix(ioHiddenMixed, d_model, seq_len, bucket, hybrid_layer_final); + + double q_mean = mean_abs_diff(cpu_q, ane_q, seq_len * (q_dim * 2)); + double q_max = max_abs_diff(cpu_q, ane_q, seq_len * (q_dim * 2)); + double query_mean = mean_abs_diff_q_half(cpu_q, ane_q, seq_len, q_dim, 0); + double query_max = max_abs_diff_q_half(cpu_q, ane_q, seq_len, q_dim, 0); + double gate_mean = mean_abs_diff_q_half(cpu_q, ane_q, seq_len, q_dim, 1); + double gate_max = max_abs_diff_q_half(cpu_q, ane_q, seq_len, q_dim, 1); + double q_half_swap_mean = mean_abs_diff_q_half_swapped(cpu_q, ane_q, seq_len, q_dim); + double q_half_swap_max = max_abs_diff_q_half_swapped(cpu_q, ane_q, seq_len, q_dim); + double k_mean = mean_abs_diff(cpu_k, ane_k, seq_len * kv_dim); + double k_max = max_abs_diff(cpu_k, ane_k, seq_len * kv_dim); + double input_rms_ms_mean = mean_abs_diff(cpu_input_rms_ms, ane_input_rms_ms, seq_len); + double input_rms_ms_max = max_abs_diff(cpu_input_rms_ms, ane_input_rms_ms, seq_len); + double input_rms_rrms_mean = mean_abs_diff(cpu_input_rms_rrms, ane_input_rms_rrms, seq_len); + double input_rms_rrms_max = max_abs_diff(cpu_input_rms_rrms, ane_input_rms_rrms, seq_len); + double input_rms_rrms_pow_micro_mean = mean_abs_diff(cpu_input_rms_rrms, ane_input_rms_rrms_pow_micro, seq_len); + double input_rms_rrms_pow_micro_max = max_abs_diff(cpu_input_rms_rrms, ane_input_rms_rrms_pow_micro, seq_len); + double input_rms_rrms_stage_vs_pow_mean = mean_abs_diff(ane_input_rms_rrms, ane_input_rms_rrms_pow_micro, seq_len); + double input_rms_rrms_stage_vs_pow_max = max_abs_diff(ane_input_rms_rrms, ane_input_rms_rrms_pow_micro, seq_len); + double input_rms_xr_mean = mean_abs_diff(cpu_input_rms_xr, ane_input_rms_xr, total); + double input_rms_xr_max = max_abs_diff(cpu_input_rms_xr, ane_input_rms_xr, total); + double input_rms_mean = mean_abs_diff(normed, ane_input_rms, total); + double input_rms_max = max_abs_diff(normed, ane_input_rms, total); + double v_mean = mean_abs_diff(cpu_v, ane_v, seq_len * kv_dim); + double v_max = max_abs_diff(cpu_v, ane_v, seq_len * kv_dim); + double v_rms_only_mean = mean_abs_diff(cpu_v, cpu_v_rms_only, seq_len * kv_dim); + double v_rms_only_max = max_abs_diff(cpu_v, cpu_v_rms_only, seq_len * kv_dim); + double v_single_mean = mean_abs_diff(cpu_v, ane_v_single, seq_len * kv_dim); + double v_single_max = max_abs_diff(cpu_v, ane_v_single, seq_len * kv_dim); + double v_linear_only_mean = mean_abs_diff(cpu_v, ane_v_linear_only, seq_len * kv_dim); + double v_linear_only_max = max_abs_diff(cpu_v, ane_v_linear_only, seq_len * kv_dim); + double v_input_fp16_cpu_mean = mean_abs_diff(cpu_v, cpu_v_input_fp16, seq_len * kv_dim); + double v_input_fp16_cpu_max = max_abs_diff(cpu_v, cpu_v_input_fp16, seq_len * kv_dim); + double v_input_fp16_ane_mean = mean_abs_diff(cpu_v_input_fp16, ane_v_linear_only_fp16input, seq_len * kv_dim); + double v_input_fp16_ane_max = max_abs_diff(cpu_v_input_fp16, ane_v_linear_only_fp16input, seq_len * kv_dim); + double v_input_fp16_ane_self_mean = mean_abs_diff(ane_v_linear_only, ane_v_linear_only_fp16input, seq_len * kv_dim); + double v_input_fp16_ane_self_max = max_abs_diff(ane_v_linear_only, ane_v_linear_only_fp16input, seq_len * kv_dim); + double v_fp16emu_cpu_mean = mean_abs_diff(cpu_v, cpu_v_linear_fp16emu, seq_len * kv_dim); + double v_fp16emu_cpu_max = max_abs_diff(cpu_v, cpu_v_linear_fp16emu, seq_len * kv_dim); + double v_fp16emu_ane_mean = mean_abs_diff(cpu_v_linear_fp16emu, ane_v_linear_only, seq_len * kv_dim); + double v_fp16emu_ane_max = max_abs_diff(cpu_v_linear_fp16emu, ane_v_linear_only, seq_len * kv_dim); + double v_fp16acc32_cpu_mean = mean_abs_diff(cpu_v, cpu_v_linear_fp16acc32, seq_len * kv_dim); + double v_fp16acc32_cpu_max = max_abs_diff(cpu_v, cpu_v_linear_fp16acc32, seq_len * kv_dim); + double v_fp16acc32_ane_mean = mean_abs_diff(cpu_v_linear_fp16acc32, ane_v_linear_only, seq_len * kv_dim); + double v_fp16acc32_ane_max = max_abs_diff(cpu_v_linear_fp16acc32, ane_v_linear_only, seq_len * kv_dim); + double v_multi_vs_single_mean = mean_abs_diff(ane_v, ane_v_single, seq_len * kv_dim); + double v_multi_vs_single_max = max_abs_diff(ane_v, ane_v_single, seq_len * kv_dim); + double kv_swap_mean = mean_abs_diff_kv_swapped(cpu_k, cpu_v, ane_k, ane_v, seq_len * kv_dim); + double kv_swap_max = max_abs_diff_kv_swapped(cpu_k, cpu_v, ane_k, ane_v, seq_len * kv_dim); + double q_token_major_mean = mean_abs_diff(cpu_q, ane_q_token_major, seq_len * (q_dim * 2)); + double q_token_major_max = max_abs_diff(cpu_q, ane_q_token_major, seq_len * (q_dim * 2)); + double k_token_major_mean = mean_abs_diff(cpu_k, ane_k_token_major, seq_len * kv_dim); + double k_token_major_max = max_abs_diff(cpu_k, ane_k_token_major, seq_len * kv_dim); + double v_token_major_mean = mean_abs_diff(cpu_v, ane_v_token_major, seq_len * kv_dim); + double v_token_major_max = max_abs_diff(cpu_v, ane_v_token_major, seq_len * kv_dim); + double attn_q_norm_mean = mean_abs_diff(cpu_attn_q_normed, ane_attn_q_normed, total_q); + double attn_q_norm_max = max_abs_diff(cpu_attn_q_normed, ane_attn_q_normed, total_q); + double attn_gate_sigmoid_mean = mean_abs_diff(cpu_attn_gate_sigmoid, ane_attn_gate_sigmoid, total_q); + double attn_gate_sigmoid_max = max_abs_diff(cpu_attn_gate_sigmoid, ane_attn_gate_sigmoid, total_q); + double attn_k_norm_mean = mean_abs_diff(cpu_attn_k_normed, ane_attn_k_normed, total_kv); + double attn_k_norm_max = max_abs_diff(cpu_attn_k_normed, ane_attn_k_normed, total_kv); + double attn_q_rope_mean = mean_abs_diff(cpu_attn_q_rope, ane_attn_q_rope, total_q); + double attn_q_rope_max = max_abs_diff(cpu_attn_q_rope, ane_attn_q_rope, total_q); + double attn_k_rope_mean = mean_abs_diff(cpu_attn_k_rope, ane_attn_k_rope, total_kv); + double attn_k_rope_max = max_abs_diff(cpu_attn_k_rope, ane_attn_k_rope, total_kv); + double attn_score_mean = mean_abs_diff(cpu_attn_scores, ane_attn_scores, total_score); + double attn_score_max = max_abs_diff(cpu_attn_scores, ane_attn_scores, total_score); + double attn_softmax_mean = mean_abs_diff(cpu_attn_probs, ane_attn_probs, total_score); + double attn_softmax_max = max_abs_diff(cpu_attn_probs, ane_attn_probs, total_score); + double attn_context_mean = mean_abs_diff(cpu_attn_context, ane_attn_context, total_q); + double attn_context_max = max_abs_diff(cpu_attn_context, ane_attn_context, total_q); + double attn_gated_context_mean = mean_abs_diff(cpu_attn_gated_context, ane_attn_gated_context, total_q); + double attn_gated_context_max = max_abs_diff(cpu_attn_gated_context, ane_attn_gated_context, total_q); + double attn_mean = mean_abs_diff(cpu_attn, ane_attn, total); + double attn_max = max_abs_diff(cpu_attn, ane_attn, total); + double attn_gate_only_mean = mean_abs_diff(cpu_attn, attn_gate_only, total); + double attn_gate_only_max = max_abs_diff(cpu_attn, attn_gate_only, total); + double attn_v_only_mean = mean_abs_diff(cpu_attn, attn_v_only, total); + double attn_v_only_max = max_abs_diff(cpu_attn, attn_v_only, total); + double attn_v_single_only_mean = mean_abs_diff(cpu_attn, attn_v_single_only, total); + double attn_v_single_only_max = max_abs_diff(cpu_attn, attn_v_single_only, total); + double attn_v_linear_only_mean = mean_abs_diff(cpu_attn, attn_v_linear_only, total); + double attn_v_linear_only_max = max_abs_diff(cpu_attn, attn_v_linear_only, total); + double attn_sigmoid_only_mean = mean_abs_diff(cpu_attn, attn_sigmoid_only, total); + double attn_sigmoid_only_max = max_abs_diff(cpu_attn, attn_sigmoid_only, total); + double attn_v_input_fp16_cpu_mean = mean_abs_diff(cpu_attn, attn_v_input_fp16_cpu, total); + double attn_v_input_fp16_cpu_max = max_abs_diff(cpu_attn, attn_v_input_fp16_cpu, total); + double attn_v_input_fp16_ane_mean = mean_abs_diff(attn_v_input_fp16_cpu, attn_v_input_fp16_ane, total); + double attn_v_input_fp16_ane_max = max_abs_diff(attn_v_input_fp16_cpu, attn_v_input_fp16_ane, total); + double attn_v_input_fp16_ane_self_mean = mean_abs_diff(attn_v_linear_only, attn_v_input_fp16_ane, total); + double attn_v_input_fp16_ane_self_max = max_abs_diff(attn_v_linear_only, attn_v_input_fp16_ane, total); + double attn_v_fp16emu_cpu_mean = mean_abs_diff(cpu_attn, attn_v_linear_fp16emu, total); + double attn_v_fp16emu_cpu_max = max_abs_diff(cpu_attn, attn_v_linear_fp16emu, total); + double attn_v_fp16emu_ane_mean = mean_abs_diff(attn_v_linear_fp16emu, attn_v_linear_only, total); + double attn_v_fp16emu_ane_max = max_abs_diff(attn_v_linear_fp16emu, attn_v_linear_only, total); + double attn_v_fp16acc32_cpu_mean = mean_abs_diff(cpu_attn, attn_v_linear_fp16acc32, total); + double attn_v_fp16acc32_cpu_max = max_abs_diff(cpu_attn, attn_v_linear_fp16acc32, total); + double attn_v_fp16acc32_ane_mean = mean_abs_diff(attn_v_linear_fp16acc32, attn_v_linear_only, total); + double attn_v_fp16acc32_ane_max = max_abs_diff(attn_v_linear_fp16acc32, attn_v_linear_only, total); + double attn_v_rms_only_mean = mean_abs_diff(cpu_attn, attn_v_rms_only, total); + double attn_v_rms_only_max = max_abs_diff(cpu_attn, attn_v_rms_only, total); + double hidden_attn_mean = mean_abs_diff(cpu_hidden_attn, ane_hidden_attn, total); + double hidden_attn_max = max_abs_diff(cpu_hidden_attn, ane_hidden_attn, total); + double ffn_rms_mean = mean_abs_diff(cpu_ffn_rms, ane_ffn_rms, total); + double ffn_rms_max = max_abs_diff(cpu_ffn_rms, ane_ffn_rms, total); + double ffn_gate_mean = mean_abs_diff(cpu_ffn_gate, ane_ffn_gate, total_ff); + double ffn_gate_max = max_abs_diff(cpu_ffn_gate, ane_ffn_gate, total_ff); + double ffn_up_mean = mean_abs_diff(cpu_ffn_up, ane_ffn_up, total_ff); + double ffn_up_max = max_abs_diff(cpu_ffn_up, ane_ffn_up, total_ff); + double ffn_up_only_mean = mean_abs_diff(cpu_ffn_up, ane_ffn_up_only, total_ff); + double ffn_up_only_max = max_abs_diff(cpu_ffn_up, ane_ffn_up_only, total_ff); + double ffn_sigmoid_micro_mean = mean_abs_diff(cpu_ffn_sigmoid, ane_ffn_sigmoid_micro, total_ff); + double ffn_sigmoid_micro_max = max_abs_diff(cpu_ffn_sigmoid, ane_ffn_sigmoid_micro, total_ff); + double ffn_sigmoid_tanh_micro_mean = mean_abs_diff(cpu_ffn_sigmoid, ane_ffn_sigmoid_tanh_micro, total_ff); + double ffn_sigmoid_tanh_micro_max = max_abs_diff(cpu_ffn_sigmoid, ane_ffn_sigmoid_tanh_micro, total_ff); + double ffn_sigmoid_fp16emu_cpu_mean = mean_abs_diff(cpu_ffn_sigmoid, cpu_ffn_sigmoid_fp16emu, total_ff); + double ffn_sigmoid_fp16emu_cpu_max = max_abs_diff(cpu_ffn_sigmoid, cpu_ffn_sigmoid_fp16emu, total_ff); + double ffn_sigmoid_fp16emu_ane_mean = mean_abs_diff(cpu_ffn_sigmoid_fp16emu, ane_ffn_sigmoid_micro, total_ff); + double ffn_sigmoid_fp16emu_ane_max = max_abs_diff(cpu_ffn_sigmoid_fp16emu, ane_ffn_sigmoid_micro, total_ff); + double ffn_silu_mean = mean_abs_diff(cpu_ffn_silu, ane_ffn_silu, total_ff); + double ffn_silu_max = max_abs_diff(cpu_ffn_silu, ane_ffn_silu, total_ff); + double ffn_silu_only_mean = mean_abs_diff(cpu_ffn_silu, ane_ffn_silu_only, total_ff); + double ffn_silu_only_max = max_abs_diff(cpu_ffn_silu, ane_ffn_silu_only, total_ff); + double ffn_silu_micro_mean = mean_abs_diff(cpu_ffn_silu, ane_ffn_silu_micro, total_ff); + double ffn_silu_micro_max = max_abs_diff(cpu_ffn_silu, ane_ffn_silu_micro, total_ff); + double ffn_silu_tanh_micro_mean = mean_abs_diff(cpu_ffn_silu, ane_ffn_silu_tanh_micro, total_ff); + double ffn_silu_tanh_micro_max = max_abs_diff(cpu_ffn_silu, ane_ffn_silu_tanh_micro, total_ff); + double ffn_silu_fp16emu_cpu_mean = mean_abs_diff(cpu_ffn_silu, cpu_ffn_silu_fp16emu, total_ff); + double ffn_silu_fp16emu_cpu_max = max_abs_diff(cpu_ffn_silu, cpu_ffn_silu_fp16emu, total_ff); + double ffn_silu_fp16emu_ane_mean = mean_abs_diff(cpu_ffn_silu_fp16emu, ane_ffn_silu_micro, total_ff); + double ffn_silu_fp16emu_ane_max = max_abs_diff(cpu_ffn_silu_fp16emu, ane_ffn_silu_micro, total_ff); + double ffn_hidden_mean = mean_abs_diff(cpu_ffn_hidden, ane_ffn_hidden, total_ff); + double ffn_hidden_max = max_abs_diff(cpu_ffn_hidden, ane_ffn_hidden, total_ff); + double ffn_hidden_only_mean = mean_abs_diff(cpu_ffn_hidden, ane_ffn_hidden_only, total_ff); + double ffn_hidden_only_max = max_abs_diff(cpu_ffn_hidden, ane_ffn_hidden_only, total_ff); + double ffn_mul_micro_mean = mean_abs_diff(cpu_ffn_hidden, ane_ffn_mul_micro, total_ff); + double ffn_mul_micro_max = max_abs_diff(cpu_ffn_hidden, ane_ffn_mul_micro, total_ff); + double ffn_hidden_fp16emu_cpu_mean = mean_abs_diff(cpu_ffn_hidden, cpu_ffn_hidden_fp16emu, total_ff); + double ffn_hidden_fp16emu_cpu_max = max_abs_diff(cpu_ffn_hidden, cpu_ffn_hidden_fp16emu, total_ff); + double ffn_hidden_fp16emu_ane_mean = mean_abs_diff(cpu_ffn_hidden_fp16emu, ane_ffn_hidden_only, total_ff); + double ffn_hidden_fp16emu_ane_max = max_abs_diff(cpu_ffn_hidden_fp16emu, ane_ffn_hidden_only, total_ff); + double ffn_down_same_input_mean = mean_abs_diff(cpu_ffn_down, ane_ffn_down_same_input, total); + double ffn_down_same_input_max = max_abs_diff(cpu_ffn_down, ane_ffn_down_same_input, total); + double ffn_same_input_mean = mean_abs_diff(cpu_ffn_final, ane_ffn_final_same_input, total); + double ffn_same_input_max = max_abs_diff(cpu_ffn_final, ane_ffn_final_same_input, total); + double full_layer_mean = mean_abs_diff(cpu_ffn_final, hybrid_layer_final, total); + double full_layer_max = max_abs_diff(cpu_ffn_final, hybrid_layer_final, total); + float cpu_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float hybrid_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float gate_only_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float sigmoid_only_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float v_only_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float cpu_ffn_final_pair_logits[2] = {0.0f, 0.0f}; + float ane_ffn_same_input_pair_logits[2] = {0.0f, 0.0f}; + float hybrid_layer_final_pair_logits[2] = {0.0f, 0.0f}; + double cpu_hidden_attn_pair_gap = 0.0; + double hybrid_hidden_attn_pair_gap = 0.0; + double gate_only_hidden_attn_pair_gap = 0.0; + double sigmoid_only_hidden_attn_pair_gap = 0.0; + double v_only_hidden_attn_pair_gap = 0.0; + double cpu_ffn_final_pair_gap = 0.0; + double ane_ffn_same_input_pair_gap = 0.0; + double hybrid_layer_final_pair_gap = 0.0; + int cpu_hidden_attn_pref_token = -1; + int hybrid_hidden_attn_pref_token = -1; + int gate_only_hidden_attn_pref_token = -1; + int sigmoid_only_hidden_attn_pref_token = -1; + int v_only_hidden_attn_pref_token = -1; + int cpu_ffn_final_pref_token = -1; + int ane_ffn_same_input_pref_token = -1; + int hybrid_layer_final_pref_token = -1; + + if (pair_enabled) { + const int pair_ids[2] = {candidate_a, candidate_b}; + const float *cpu_hidden_attn_last = cpu_hidden_attn + (seq_len - 1) * d_model; + const float *ane_hidden_attn_last = ane_hidden_attn + (seq_len - 1) * d_model; + const float *gate_only_hidden_attn_last = gate_only_hidden_attn + (seq_len - 1) * d_model; + const float *sigmoid_only_hidden_attn_last = sigmoid_only_hidden_attn + (seq_len - 1) * d_model; + const float *v_only_hidden_attn_last = v_only_hidden_attn + (seq_len - 1) * d_model; + const float *cpu_ffn_final_last = cpu_ffn_final + (seq_len - 1) * d_model; + const float *ane_ffn_final_same_input_last = ane_ffn_final_same_input + (seq_len - 1) * d_model; + const float *hybrid_layer_final_last = hybrid_layer_final + (seq_len - 1) * d_model; + + orion_qwen_cpu_rmsnorm(cpu_hidden_attn_last, final_norm, d_model, 1e-6f, cpu_stage_last); + orion_qwen_cpu_rmsnorm(ane_hidden_attn_last, final_norm, d_model, 1e-6f, hybrid_stage_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_stage_last, d_model, pair_ids, 2, cpu_hidden_attn_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_stage_last, d_model, pair_ids, 2, hybrid_hidden_attn_pair_logits)) { + fprintf(stderr, "FAIL: pair logits failed at hidden_attn stage\n"); + goto fail; + } + orion_qwen_cpu_rmsnorm(gate_only_hidden_attn_last, final_norm, d_model, 1e-6f, cpu_stage_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_stage_last, d_model, pair_ids, 2, gate_only_hidden_attn_pair_logits)) { + fprintf(stderr, "FAIL: pair logits failed at gate-only hidden_attn stage\n"); + goto fail; + } + orion_qwen_cpu_rmsnorm(sigmoid_only_hidden_attn_last, final_norm, d_model, 1e-6f, cpu_stage_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_stage_last, d_model, pair_ids, 2, sigmoid_only_hidden_attn_pair_logits)) { + fprintf(stderr, "FAIL: pair logits failed at sigmoid-only hidden_attn stage\n"); + goto fail; + } + orion_qwen_cpu_rmsnorm(v_only_hidden_attn_last, final_norm, d_model, 1e-6f, cpu_stage_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_stage_last, d_model, pair_ids, 2, v_only_hidden_attn_pair_logits)) { + fprintf(stderr, "FAIL: pair logits failed at v-only hidden_attn stage\n"); + goto fail; + } + + orion_qwen_cpu_rmsnorm(cpu_ffn_final_last, final_norm, d_model, 1e-6f, cpu_stage_last); + orion_qwen_cpu_rmsnorm(ane_ffn_final_same_input_last, final_norm, d_model, 1e-6f, hybrid_stage_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_stage_last, d_model, pair_ids, 2, cpu_ffn_final_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_stage_last, d_model, pair_ids, 2, ane_ffn_same_input_pair_logits)) { + fprintf(stderr, "FAIL: pair logits failed at ffn same-input stage\n"); + goto fail; + } + + orion_qwen_cpu_rmsnorm(hybrid_layer_final_last, final_norm, d_model, 1e-6f, hybrid_stage_last); + if (!selected_token_logits(blob_dir, lm_head_name, hybrid_stage_last, d_model, pair_ids, 2, hybrid_layer_final_pair_logits)) { + fprintf(stderr, "FAIL: pair logits failed at hybrid final stage\n"); + goto fail; + } + + cpu_hidden_attn_pair_gap = (double)cpu_hidden_attn_pair_logits[0] - (double)cpu_hidden_attn_pair_logits[1]; + hybrid_hidden_attn_pair_gap = (double)hybrid_hidden_attn_pair_logits[0] - (double)hybrid_hidden_attn_pair_logits[1]; + gate_only_hidden_attn_pair_gap = (double)gate_only_hidden_attn_pair_logits[0] - (double)gate_only_hidden_attn_pair_logits[1]; + sigmoid_only_hidden_attn_pair_gap = (double)sigmoid_only_hidden_attn_pair_logits[0] - (double)sigmoid_only_hidden_attn_pair_logits[1]; + v_only_hidden_attn_pair_gap = (double)v_only_hidden_attn_pair_logits[0] - (double)v_only_hidden_attn_pair_logits[1]; + cpu_ffn_final_pair_gap = (double)cpu_ffn_final_pair_logits[0] - (double)cpu_ffn_final_pair_logits[1]; + ane_ffn_same_input_pair_gap = (double)ane_ffn_same_input_pair_logits[0] - (double)ane_ffn_same_input_pair_logits[1]; + hybrid_layer_final_pair_gap = (double)hybrid_layer_final_pair_logits[0] - (double)hybrid_layer_final_pair_logits[1]; + cpu_hidden_attn_pref_token = (cpu_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + hybrid_hidden_attn_pref_token = (hybrid_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + gate_only_hidden_attn_pref_token = (gate_only_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + sigmoid_only_hidden_attn_pref_token = (sigmoid_only_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + v_only_hidden_attn_pref_token = (v_only_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + cpu_ffn_final_pref_token = (cpu_ffn_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + ane_ffn_same_input_pref_token = (ane_ffn_same_input_pair_gap >= 0.0) ? candidate_a : candidate_b; + hybrid_layer_final_pref_token = (hybrid_layer_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + } + + printf("PASS: qwen35 9b bridge stage diff trace\n"); + printf(" prompt=%s\n", prompt); + printf(" prompt_len=%d\n", seq_len); + printf(" target_layer=%d\n", layer); + printf(" candidate_a=%d\n", candidate_a); + printf(" candidate_b=%d\n", candidate_b); + printf(" q_proj_source=%s\n", q_proj_source_label(q_proj_uses_cpu)); + printf(" q_query_source=%s\n", q_query_source_label(q_query_uses_cpu)); + printf(" q_gate_source=%s\n", q_gate_source_label(q_gate_uses_cpu)); + printf(" v_proj_source=%s\n", v_proj_source_label(v_proj_uses_cpu, v_proj_cpu_channel_count)); + printf(" v_proj_cpu_channel_count=%d\n", v_proj_cpu_channel_count); + printf(" trace_dim_count=%d\n", trace_dim_count); + printf(" trace_channel_count=%d\n", trace_channel_count); + for (int i = 0; i < trace_dim_count; i++) { + printf(" trace_dim_config_rank=%d dim=%d\n", i + 1, trace_dims[i]); + } + for (int i = 0; i < trace_channel_count; i++) { + printf(" trace_channel_config_rank=%d channel=%d\n", i + 1, trace_channels[i]); + } + printf(" q_proj_mean_abs_diff=%.6f\n", q_mean); + printf(" q_proj_max_abs_diff=%.6f\n", q_max); + printf(" q_query_mean_abs_diff=%.6f\n", query_mean); + printf(" q_query_max_abs_diff=%.6f\n", query_max); + printf(" q_gate_mean_abs_diff=%.6f\n", gate_mean); + printf(" q_gate_max_abs_diff=%.6f\n", gate_max); + printf(" q_half_swap_mean_abs_diff=%.6f\n", q_half_swap_mean); + printf(" q_half_swap_max_abs_diff=%.6f\n", q_half_swap_max); + printf(" k_proj_mean_abs_diff=%.6f\n", k_mean); + printf(" k_proj_max_abs_diff=%.6f\n", k_max); + printf(" input_rms_ms_mean_abs_diff=%.6f\n", input_rms_ms_mean); + printf(" input_rms_ms_max_abs_diff=%.6f\n", input_rms_ms_max); + printf(" input_rms_rrms_mean_abs_diff=%.6f\n", input_rms_rrms_mean); + printf(" input_rms_rrms_max_abs_diff=%.6f\n", input_rms_rrms_max); + printf(" input_rms_rrms_pow_micro_mean_abs_diff=%.6f\n", input_rms_rrms_pow_micro_mean); + printf(" input_rms_rrms_pow_micro_max_abs_diff=%.6f\n", input_rms_rrms_pow_micro_max); + printf(" input_rms_rrms_stage_vs_pow_mean_abs_diff=%.6f\n", input_rms_rrms_stage_vs_pow_mean); + printf(" input_rms_rrms_stage_vs_pow_max_abs_diff=%.6f\n", input_rms_rrms_stage_vs_pow_max); + printf(" input_rms_xr_mean_abs_diff=%.6f\n", input_rms_xr_mean); + printf(" input_rms_xr_max_abs_diff=%.6f\n", input_rms_xr_max); + printf(" input_rms_mean_abs_diff=%.6f\n", input_rms_mean); + printf(" input_rms_max_abs_diff=%.6f\n", input_rms_max); + printf(" v_proj_mean_abs_diff=%.6f\n", v_mean); + printf(" v_proj_max_abs_diff=%.6f\n", v_max); + printf(" v_proj_rms_only_mean_abs_diff=%.6f\n", v_rms_only_mean); + printf(" v_proj_rms_only_max_abs_diff=%.6f\n", v_rms_only_max); + printf(" v_proj_single_mean_abs_diff=%.6f\n", v_single_mean); + printf(" v_proj_single_max_abs_diff=%.6f\n", v_single_max); + printf(" v_proj_linear_only_mean_abs_diff=%.6f\n", v_linear_only_mean); + printf(" v_proj_linear_only_max_abs_diff=%.6f\n", v_linear_only_max); + printf(" v_proj_input_fp16_cpu_mean_abs_diff=%.6f\n", v_input_fp16_cpu_mean); + printf(" v_proj_input_fp16_cpu_max_abs_diff=%.6f\n", v_input_fp16_cpu_max); + printf(" v_proj_input_fp16_ane_mean_abs_diff=%.6f\n", v_input_fp16_ane_mean); + printf(" v_proj_input_fp16_ane_max_abs_diff=%.6f\n", v_input_fp16_ane_max); + printf(" v_proj_input_fp16_ane_self_mean_abs_diff=%.6f\n", v_input_fp16_ane_self_mean); + printf(" v_proj_input_fp16_ane_self_max_abs_diff=%.6f\n", v_input_fp16_ane_self_max); + printf(" v_proj_fp16emu_cpu_mean_abs_diff=%.6f\n", v_fp16emu_cpu_mean); + printf(" v_proj_fp16emu_cpu_max_abs_diff=%.6f\n", v_fp16emu_cpu_max); + printf(" v_proj_fp16emu_ane_mean_abs_diff=%.6f\n", v_fp16emu_ane_mean); + printf(" v_proj_fp16emu_ane_max_abs_diff=%.6f\n", v_fp16emu_ane_max); + printf(" v_proj_fp16acc32_cpu_mean_abs_diff=%.6f\n", v_fp16acc32_cpu_mean); + printf(" v_proj_fp16acc32_cpu_max_abs_diff=%.6f\n", v_fp16acc32_cpu_max); + printf(" v_proj_fp16acc32_ane_mean_abs_diff=%.6f\n", v_fp16acc32_ane_mean); + printf(" v_proj_fp16acc32_ane_max_abs_diff=%.6f\n", v_fp16acc32_ane_max); + printf(" v_proj_multi_vs_single_mean_abs_diff=%.6f\n", v_multi_vs_single_mean); + printf(" v_proj_multi_vs_single_max_abs_diff=%.6f\n", v_multi_vs_single_max); + printf(" v_linear_only_tmp_dir=%s\n", v_linear_only_tmp_dir ? v_linear_only_tmp_dir.UTF8String : ""); + printf(" v_linear_only_source_blob_path=%s\n", v_linear_only_source_blob_path ? v_linear_only_source_blob_path.UTF8String : ""); + printf(" v_linear_only_data_path=%s\n", v_linear_only_data_path ? v_linear_only_data_path.UTF8String : ""); + printf(" v_linear_only_weight_path=%s\n", v_linear_only_weight_path ? v_linear_only_weight_path.UTF8String : ""); + printf(" v_linear_only_source_blob_bytes=%lld\n", v_linear_only_data_vs_source.rhs_size); + printf(" v_linear_only_source_payload_bytes=%lld\n", v_linear_only_data_vs_source_payload.rhs_size); + printf(" v_linear_only_data_bytes=%lld\n", v_linear_only_data_vs_source.lhs_size); + printf(" v_linear_only_weight_bytes=%lld\n", v_linear_only_weight_vs_source.lhs_size); + printf(" v_linear_only_data_matches_source_blob=%d\n", v_linear_only_data_vs_source.equal); + printf(" v_linear_only_weight_matches_source_blob=%d\n", v_linear_only_weight_vs_source.equal); + printf(" v_linear_only_data_matches_weight_blob=%d\n", v_linear_only_data_vs_weight.equal); + printf(" v_linear_only_data_matches_source_payload=%d\n", v_linear_only_data_vs_source_payload.equal); + printf(" v_linear_only_data_vs_source_first_diff=%lld\n", v_linear_only_data_vs_source.first_diff_offset); + printf(" v_linear_only_weight_vs_source_first_diff=%lld\n", v_linear_only_weight_vs_source.first_diff_offset); + printf(" v_linear_only_data_vs_weight_first_diff=%lld\n", v_linear_only_data_vs_weight.first_diff_offset); + printf(" v_linear_only_data_vs_source_payload_first_diff=%lld\n", v_linear_only_data_vs_source_payload.first_diff_offset); + printf(" kv_swap_mean_abs_diff=%.6f\n", kv_swap_mean); + printf(" kv_swap_max_abs_diff=%.6f\n", kv_swap_max); + printf(" q_token_major_mean_abs_diff=%.6f\n", q_token_major_mean); + printf(" q_token_major_max_abs_diff=%.6f\n", q_token_major_max); + printf(" k_token_major_mean_abs_diff=%.6f\n", k_token_major_mean); + printf(" k_token_major_max_abs_diff=%.6f\n", k_token_major_max); + printf(" v_token_major_mean_abs_diff=%.6f\n", v_token_major_mean); + printf(" v_token_major_max_abs_diff=%.6f\n", v_token_major_max); + printf(" attn_q_norm_mean_abs_diff=%.6f\n", attn_q_norm_mean); + printf(" attn_q_norm_max_abs_diff=%.6f\n", attn_q_norm_max); + printf(" attn_gate_sigmoid_mean_abs_diff=%.6f\n", attn_gate_sigmoid_mean); + printf(" attn_gate_sigmoid_max_abs_diff=%.6f\n", attn_gate_sigmoid_max); + printf(" attn_k_norm_mean_abs_diff=%.6f\n", attn_k_norm_mean); + printf(" attn_k_norm_max_abs_diff=%.6f\n", attn_k_norm_max); + printf(" attn_q_rope_mean_abs_diff=%.6f\n", attn_q_rope_mean); + printf(" attn_q_rope_max_abs_diff=%.6f\n", attn_q_rope_max); + printf(" attn_k_rope_mean_abs_diff=%.6f\n", attn_k_rope_mean); + printf(" attn_k_rope_max_abs_diff=%.6f\n", attn_k_rope_max); + printf(" attn_score_mean_abs_diff=%.6f\n", attn_score_mean); + printf(" attn_score_max_abs_diff=%.6f\n", attn_score_max); + printf(" attn_softmax_mean_abs_diff=%.6f\n", attn_softmax_mean); + printf(" attn_softmax_max_abs_diff=%.6f\n", attn_softmax_max); + printf(" attn_context_mean_abs_diff=%.6f\n", attn_context_mean); + printf(" attn_context_max_abs_diff=%.6f\n", attn_context_max); + printf(" attn_gated_context_mean_abs_diff=%.6f\n", attn_gated_context_mean); + printf(" attn_gated_context_max_abs_diff=%.6f\n", attn_gated_context_max); + printf(" attn_out_mean_abs_diff=%.6f\n", attn_mean); + printf(" attn_out_max_abs_diff=%.6f\n", attn_max); + printf(" attn_gate_only_mean_abs_diff=%.6f\n", attn_gate_only_mean); + printf(" attn_gate_only_max_abs_diff=%.6f\n", attn_gate_only_max); + printf(" attn_v_only_mean_abs_diff=%.6f\n", attn_v_only_mean); + printf(" attn_v_only_max_abs_diff=%.6f\n", attn_v_only_max); + printf(" attn_v_single_only_mean_abs_diff=%.6f\n", attn_v_single_only_mean); + printf(" attn_v_single_only_max_abs_diff=%.6f\n", attn_v_single_only_max); + printf(" attn_v_linear_only_mean_abs_diff=%.6f\n", attn_v_linear_only_mean); + printf(" attn_v_linear_only_max_abs_diff=%.6f\n", attn_v_linear_only_max); + printf(" attn_sigmoid_only_mean_abs_diff=%.6f\n", attn_sigmoid_only_mean); + printf(" attn_sigmoid_only_max_abs_diff=%.6f\n", attn_sigmoid_only_max); + printf(" attn_v_input_fp16_cpu_mean_abs_diff=%.6f\n", attn_v_input_fp16_cpu_mean); + printf(" attn_v_input_fp16_cpu_max_abs_diff=%.6f\n", attn_v_input_fp16_cpu_max); + printf(" attn_v_input_fp16_ane_mean_abs_diff=%.6f\n", attn_v_input_fp16_ane_mean); + printf(" attn_v_input_fp16_ane_max_abs_diff=%.6f\n", attn_v_input_fp16_ane_max); + printf(" attn_v_input_fp16_ane_self_mean_abs_diff=%.6f\n", attn_v_input_fp16_ane_self_mean); + printf(" attn_v_input_fp16_ane_self_max_abs_diff=%.6f\n", attn_v_input_fp16_ane_self_max); + printf(" attn_v_fp16emu_cpu_mean_abs_diff=%.6f\n", attn_v_fp16emu_cpu_mean); + printf(" attn_v_fp16emu_cpu_max_abs_diff=%.6f\n", attn_v_fp16emu_cpu_max); + printf(" attn_v_fp16emu_ane_mean_abs_diff=%.6f\n", attn_v_fp16emu_ane_mean); + printf(" attn_v_fp16emu_ane_max_abs_diff=%.6f\n", attn_v_fp16emu_ane_max); + printf(" attn_v_fp16acc32_cpu_mean_abs_diff=%.6f\n", attn_v_fp16acc32_cpu_mean); + printf(" attn_v_fp16acc32_cpu_max_abs_diff=%.6f\n", attn_v_fp16acc32_cpu_max); + printf(" attn_v_fp16acc32_ane_mean_abs_diff=%.6f\n", attn_v_fp16acc32_ane_mean); + printf(" attn_v_fp16acc32_ane_max_abs_diff=%.6f\n", attn_v_fp16acc32_ane_max); + printf(" attn_v_rms_only_mean_abs_diff=%.6f\n", attn_v_rms_only_mean); + printf(" attn_v_rms_only_max_abs_diff=%.6f\n", attn_v_rms_only_max); + printf(" hidden_attn_mean_abs_diff=%.6f\n", hidden_attn_mean); + printf(" hidden_attn_max_abs_diff=%.6f\n", hidden_attn_max); + printf(" ffn_rms_mean_abs_diff=%.6f\n", ffn_rms_mean); + printf(" ffn_rms_max_abs_diff=%.6f\n", ffn_rms_max); + printf(" ffn_gate_mean_abs_diff=%.6f\n", ffn_gate_mean); + printf(" ffn_gate_max_abs_diff=%.6f\n", ffn_gate_max); + printf(" ffn_up_mean_abs_diff=%.6f\n", ffn_up_mean); + printf(" ffn_up_max_abs_diff=%.6f\n", ffn_up_max); + printf(" ffn_up_only_mean_abs_diff=%.6f\n", ffn_up_only_mean); + printf(" ffn_up_only_max_abs_diff=%.6f\n", ffn_up_only_max); + printf(" ffn_sigmoid_micro_mean_abs_diff=%.6f\n", ffn_sigmoid_micro_mean); + printf(" ffn_sigmoid_micro_max_abs_diff=%.6f\n", ffn_sigmoid_micro_max); + printf(" ffn_sigmoid_tanh_micro_mean_abs_diff=%.6f\n", ffn_sigmoid_tanh_micro_mean); + printf(" ffn_sigmoid_tanh_micro_max_abs_diff=%.6f\n", ffn_sigmoid_tanh_micro_max); + printf(" ffn_sigmoid_fp16emu_cpu_mean_abs_diff=%.6f\n", ffn_sigmoid_fp16emu_cpu_mean); + printf(" ffn_sigmoid_fp16emu_cpu_max_abs_diff=%.6f\n", ffn_sigmoid_fp16emu_cpu_max); + printf(" ffn_sigmoid_fp16emu_ane_mean_abs_diff=%.6f\n", ffn_sigmoid_fp16emu_ane_mean); + printf(" ffn_sigmoid_fp16emu_ane_max_abs_diff=%.6f\n", ffn_sigmoid_fp16emu_ane_max); + printf(" ffn_silu_mean_abs_diff=%.6f\n", ffn_silu_mean); + printf(" ffn_silu_max_abs_diff=%.6f\n", ffn_silu_max); + printf(" ffn_silu_only_mean_abs_diff=%.6f\n", ffn_silu_only_mean); + printf(" ffn_silu_only_max_abs_diff=%.6f\n", ffn_silu_only_max); + printf(" ffn_silu_micro_mean_abs_diff=%.6f\n", ffn_silu_micro_mean); + printf(" ffn_silu_micro_max_abs_diff=%.6f\n", ffn_silu_micro_max); + printf(" ffn_silu_tanh_micro_mean_abs_diff=%.6f\n", ffn_silu_tanh_micro_mean); + printf(" ffn_silu_tanh_micro_max_abs_diff=%.6f\n", ffn_silu_tanh_micro_max); + printf(" ffn_silu_fp16emu_cpu_mean_abs_diff=%.6f\n", ffn_silu_fp16emu_cpu_mean); + printf(" ffn_silu_fp16emu_cpu_max_abs_diff=%.6f\n", ffn_silu_fp16emu_cpu_max); + printf(" ffn_silu_fp16emu_ane_mean_abs_diff=%.6f\n", ffn_silu_fp16emu_ane_mean); + printf(" ffn_silu_fp16emu_ane_max_abs_diff=%.6f\n", ffn_silu_fp16emu_ane_max); + printf(" ffn_hidden_mean_abs_diff=%.6f\n", ffn_hidden_mean); + printf(" ffn_hidden_max_abs_diff=%.6f\n", ffn_hidden_max); + printf(" ffn_hidden_only_mean_abs_diff=%.6f\n", ffn_hidden_only_mean); + printf(" ffn_hidden_only_max_abs_diff=%.6f\n", ffn_hidden_only_max); + printf(" ffn_mul_micro_mean_abs_diff=%.6f\n", ffn_mul_micro_mean); + printf(" ffn_mul_micro_max_abs_diff=%.6f\n", ffn_mul_micro_max); + printf(" ffn_hidden_fp16emu_cpu_mean_abs_diff=%.6f\n", ffn_hidden_fp16emu_cpu_mean); + printf(" ffn_hidden_fp16emu_cpu_max_abs_diff=%.6f\n", ffn_hidden_fp16emu_cpu_max); + printf(" ffn_hidden_fp16emu_ane_mean_abs_diff=%.6f\n", ffn_hidden_fp16emu_ane_mean); + printf(" ffn_hidden_fp16emu_ane_max_abs_diff=%.6f\n", ffn_hidden_fp16emu_ane_max); + printf(" ffn_down_same_input_mean_abs_diff=%.6f\n", ffn_down_same_input_mean); + printf(" ffn_down_same_input_max_abs_diff=%.6f\n", ffn_down_same_input_max); + printf(" ffn_same_input_mean_abs_diff=%.6f\n", ffn_same_input_mean); + printf(" ffn_same_input_max_abs_diff=%.6f\n", ffn_same_input_max); + printf(" full_layer_mean_abs_diff=%.6f\n", full_layer_mean); + printf(" full_layer_max_abs_diff=%.6f\n", full_layer_max); + if (pair_enabled) { + printf(" cpu_hidden_attn_pair_gap=%.6f\n", cpu_hidden_attn_pair_gap); + printf(" hybrid_hidden_attn_pair_gap=%.6f\n", hybrid_hidden_attn_pair_gap); + printf(" gate_only_hidden_attn_pair_gap=%.6f\n", gate_only_hidden_attn_pair_gap); + printf(" sigmoid_only_hidden_attn_pair_gap=%.6f\n", sigmoid_only_hidden_attn_pair_gap); + printf(" v_only_hidden_attn_pair_gap=%.6f\n", v_only_hidden_attn_pair_gap); + printf(" cpu_hidden_attn_pref_token=%d\n", cpu_hidden_attn_pref_token); + printf(" hybrid_hidden_attn_pref_token=%d\n", hybrid_hidden_attn_pref_token); + printf(" gate_only_hidden_attn_pref_token=%d\n", gate_only_hidden_attn_pref_token); + printf(" sigmoid_only_hidden_attn_pref_token=%d\n", sigmoid_only_hidden_attn_pref_token); + printf(" v_only_hidden_attn_pref_token=%d\n", v_only_hidden_attn_pref_token); + printf(" cpu_ffn_final_pair_gap=%.6f\n", cpu_ffn_final_pair_gap); + printf(" ane_ffn_same_input_pair_gap=%.6f\n", ane_ffn_same_input_pair_gap); + printf(" hybrid_layer_final_pair_gap=%.6f\n", hybrid_layer_final_pair_gap); + printf(" cpu_ffn_final_pref_token=%d\n", cpu_ffn_final_pref_token); + printf(" ane_ffn_same_input_pref_token=%d\n", ane_ffn_same_input_pref_token); + printf(" hybrid_layer_final_pref_token=%d\n", hybrid_layer_final_pref_token); + } + print_trace_stage_rows("input_rms_xr", cpu_input_rms_xr, ane_input_rms_xr, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("input_rms_out", normed, ane_input_rms, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("attn_out", cpu_attn, ane_attn, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("hidden_attn", cpu_hidden_attn, ane_hidden_attn, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("ffn_rms", cpu_ffn_rms, ane_ffn_rms, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("ffn_down_same_input", cpu_ffn_down, ane_ffn_down_same_input, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("ffn_final_same_input", cpu_ffn_final, ane_ffn_final_same_input, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("full_layer_final", cpu_ffn_final, hybrid_layer_final, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("attn_gate_only", cpu_attn, attn_gate_only, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("attn_v_only", cpu_attn, attn_v_only, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("attn_v_single_only", cpu_attn, attn_v_single_only, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("attn_v_linear_only", cpu_attn, attn_v_linear_only, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_stage_rows("attn_v_rms_only", cpu_attn, attn_v_rms_only, seq_len, d_model, trace_dims, trace_dim_count); + print_trace_attn_attr_rows(cpu_attn_gated_context, ane_attn_gated_context, o_proj, + seq_len, q_dim, head_dim, trace_dims, trace_dim_count); + print_trace_attn_channel_rows(cpu_q, ane_q, cpu_v, cpu_v_rms_only, ane_v, ane_v_single, ane_v_linear_only, + cpu_attn_q_normed, ane_attn_q_normed, + cpu_attn_gate_sigmoid, ane_attn_gate_sigmoid, + cpu_attn_q_rope, ane_attn_q_rope, + cpu_attn_context, ane_attn_context, + cpu_attn_gated_context, ane_attn_gated_context, + seq_len, q_dim, kv_dim, head_dim, q_per_kv, + trace_channels, trace_channel_count); + print_trace_v_neighbor_rows("v_proj", cpu_v, ane_v, seq_len, kv_dim, head_dim, q_per_kv, trace_channels, trace_channel_count); + print_trace_v_neighbor_rows("v_proj_linear_only", cpu_v, ane_v_linear_only, seq_len, kv_dim, head_dim, q_per_kv, trace_channels, trace_channel_count); + + if (ioIn) CFRelease(ioIn); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioVSingle) CFRelease(ioVSingle); + if (ioNormedIn) CFRelease(ioNormedIn); + if (ioNormedFp16In) CFRelease(ioNormedFp16In); + if (ioInputRmsOnly) CFRelease(ioInputRmsOnly); + if (ioInputRmsMs) CFRelease(ioInputRmsMs); + if (ioInputRmsRrms) CFRelease(ioInputRmsRrms); + if (ioInputRmsMseIn) CFRelease(ioInputRmsMseIn); + if (ioInputRmsRrmsPowMicro) CFRelease(ioInputRmsRrmsPowMicro); + if (ioInputRmsXr) CFRelease(ioInputRmsXr); + if (ioVLinearOnly) CFRelease(ioVLinearOnly); + if (ioVLinearOnlyFp16In) CFRelease(ioVLinearOnlyFp16In); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioFfnRms) CFRelease(ioFfnRms); + if (ioFfnGate) CFRelease(ioFfnGate); + if (ioFfnUp) CFRelease(ioFfnUp); + if (ioFfnUpOnly) CFRelease(ioFfnUpOnly); + if (ioFfnSigmoidMicroIn) CFRelease(ioFfnSigmoidMicroIn); + if (ioFfnSigmoidMicroOut) CFRelease(ioFfnSigmoidMicroOut); + if (ioFfnSigmoidTanhMicroOut) CFRelease(ioFfnSigmoidTanhMicroOut); + if (ioFfnSilu) CFRelease(ioFfnSilu); + if (ioFfnSiluOnly) CFRelease(ioFfnSiluOnly); + if (ioFfnSiluMicroIn) CFRelease(ioFfnSiluMicroIn); + if (ioFfnSiluMicroOut) CFRelease(ioFfnSiluMicroOut); + if (ioFfnSiluTanhMicroOut) CFRelease(ioFfnSiluTanhMicroOut); + if (ioFfnHidden) CFRelease(ioFfnHidden); + if (ioFfnHiddenOnly) CFRelease(ioFfnHiddenOnly); + if (ioFfnMulMicroA) CFRelease(ioFfnMulMicroA); + if (ioFfnMulMicroB) CFRelease(ioFfnMulMicroB); + if (ioFfnMulMicroOut) CFRelease(ioFfnMulMicroOut); + if (ioFfnDownIn) CFRelease(ioFfnDownIn); + if (ioFfnDown) CFRelease(ioFfnDown); + if (ioHidden) CFRelease(ioHidden); + if (ioFfnInMixed) CFRelease(ioFfnInMixed); + if (ioHiddenMixed) CFRelease(ioHiddenMixed); + bridge_release(&bridge); + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(gate_proj); free(up_proj); free(down_proj); + free(final_norm); free(cpu_stage_last); free(hybrid_stage_last); + free(hidden); free(hidden_next); free(normed); free(normed_fp16input); free(cpu_q); free(cpu_k); free(cpu_v); free(cpu_v_rms_only); free(cpu_v_linear_fp16emu); free(cpu_v_linear_fp16acc32); free(cpu_v_input_fp16); + free(cpu_input_rms_ms); free(cpu_input_rms_mse); free(cpu_input_rms_rrms); free(cpu_input_rms_xr); + free(ane_q); free(ane_k); free(ane_v); free(ane_v_single); free(ane_input_rms); + free(ane_input_rms_ms); free(ane_input_rms_rrms); free(ane_input_rms_rrms_pow_micro); free(ane_input_rms_xr); + free(ane_v_linear_only); free(ane_v_linear_only_fp16input); + free(ane_q_token_major); free(ane_k_token_major); free(ane_v_token_major); + free(cpu_attn); free(ane_attn); free(attn_gate_only); free(attn_v_only); free(attn_v_single_only); + free(attn_v_linear_only); free(attn_v_linear_fp16emu); free(attn_v_linear_fp16acc32); free(attn_v_input_fp16_cpu); free(attn_v_input_fp16_ane); free(attn_v_rms_only); free(sigmoid_only_gated_context); free(attn_sigmoid_only); free(cpu_hidden_attn); free(ane_hidden_attn); free(gate_only_hidden_attn); free(sigmoid_only_hidden_attn); free(v_only_hidden_attn); + free(cpu_attn_q_normed); free(ane_attn_q_normed); free(cpu_attn_gate_sigmoid); free(ane_attn_gate_sigmoid); + free(cpu_attn_k_normed); free(ane_attn_k_normed); free(cpu_attn_q_rope); free(ane_attn_q_rope); + free(cpu_attn_k_rope); free(ane_attn_k_rope); free(cpu_attn_scores); free(ane_attn_scores); + free(cpu_attn_probs); free(ane_attn_probs); free(cpu_attn_context); free(ane_attn_context); + free(cpu_attn_gated_context); free(ane_attn_gated_context); + free(cpu_ffn_rms); free(ane_ffn_rms); free(cpu_ffn_gate); free(ane_ffn_gate); + free(cpu_ffn_up); free(ane_ffn_up); free(ane_ffn_up_only); + free(cpu_ffn_sigmoid); free(cpu_ffn_sigmoid_fp16emu); free(ane_ffn_sigmoid_micro); free(ane_ffn_sigmoid_tanh_micro); + free(cpu_ffn_silu); free(cpu_ffn_silu_fp16emu); free(ane_ffn_silu); free(ane_ffn_silu_only); free(ane_ffn_silu_micro); free(ane_ffn_silu_tanh_micro); + free(cpu_ffn_hidden); free(cpu_ffn_hidden_fp16emu); free(ane_ffn_hidden); free(ane_ffn_hidden_only); free(ane_ffn_mul_micro); + free(cpu_ffn_down); free(ane_ffn_down_same_input); + free(cpu_ffn_final); free(ane_ffn_final_same_input); free(hybrid_layer_final); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 0; + +fail: + if (ioIn) CFRelease(ioIn); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioVSingle) CFRelease(ioVSingle); + if (ioNormedIn) CFRelease(ioNormedIn); + if (ioNormedFp16In) CFRelease(ioNormedFp16In); + if (ioInputRmsOnly) CFRelease(ioInputRmsOnly); + if (ioInputRmsMs) CFRelease(ioInputRmsMs); + if (ioInputRmsRrms) CFRelease(ioInputRmsRrms); + if (ioInputRmsMseIn) CFRelease(ioInputRmsMseIn); + if (ioInputRmsRrmsPowMicro) CFRelease(ioInputRmsRrmsPowMicro); + if (ioInputRmsXr) CFRelease(ioInputRmsXr); + if (ioVLinearOnly) CFRelease(ioVLinearOnly); + if (ioVLinearOnlyFp16In) CFRelease(ioVLinearOnlyFp16In); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioFfnRms) CFRelease(ioFfnRms); + if (ioFfnGate) CFRelease(ioFfnGate); + if (ioFfnUp) CFRelease(ioFfnUp); + if (ioFfnUpOnly) CFRelease(ioFfnUpOnly); + if (ioFfnSigmoidMicroIn) CFRelease(ioFfnSigmoidMicroIn); + if (ioFfnSigmoidMicroOut) CFRelease(ioFfnSigmoidMicroOut); + if (ioFfnSigmoidTanhMicroOut) CFRelease(ioFfnSigmoidTanhMicroOut); + if (ioFfnSilu) CFRelease(ioFfnSilu); + if (ioFfnSiluOnly) CFRelease(ioFfnSiluOnly); + if (ioFfnSiluMicroIn) CFRelease(ioFfnSiluMicroIn); + if (ioFfnSiluMicroOut) CFRelease(ioFfnSiluMicroOut); + if (ioFfnSiluTanhMicroOut) CFRelease(ioFfnSiluTanhMicroOut); + if (ioFfnHidden) CFRelease(ioFfnHidden); + if (ioFfnHiddenOnly) CFRelease(ioFfnHiddenOnly); + if (ioFfnMulMicroA) CFRelease(ioFfnMulMicroA); + if (ioFfnMulMicroB) CFRelease(ioFfnMulMicroB); + if (ioFfnMulMicroOut) CFRelease(ioFfnMulMicroOut); + if (ioFfnDownIn) CFRelease(ioFfnDownIn); + if (ioFfnDown) CFRelease(ioFfnDown); + if (ioHidden) CFRelease(ioHidden); + if (ioFfnInMixed) CFRelease(ioFfnInMixed); + if (ioHiddenMixed) CFRelease(ioHiddenMixed); + bridge_release(&bridge); + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(gate_proj); free(up_proj); free(down_proj); + free(final_norm); free(cpu_stage_last); free(hybrid_stage_last); + free(hidden); free(hidden_next); free(normed); free(normed_fp16input); free(cpu_q); free(cpu_k); free(cpu_v); free(cpu_v_rms_only); free(cpu_v_linear_fp16emu); free(cpu_v_linear_fp16acc32); free(cpu_v_input_fp16); + free(cpu_input_rms_ms); free(cpu_input_rms_mse); free(cpu_input_rms_rrms); free(cpu_input_rms_xr); + free(ane_q); free(ane_k); free(ane_v); free(ane_v_single); free(ane_input_rms); + free(ane_input_rms_ms); free(ane_input_rms_rrms); free(ane_input_rms_rrms_pow_micro); free(ane_input_rms_xr); + free(ane_v_linear_only); free(ane_v_linear_only_fp16input); + free(ane_q_token_major); free(ane_k_token_major); free(ane_v_token_major); + free(cpu_attn); free(ane_attn); free(attn_gate_only); free(attn_v_only); free(attn_v_single_only); + free(attn_v_linear_only); free(attn_v_linear_fp16emu); free(attn_v_linear_fp16acc32); free(attn_v_input_fp16_cpu); free(attn_v_input_fp16_ane); free(attn_v_rms_only); free(sigmoid_only_gated_context); free(attn_sigmoid_only); free(cpu_hidden_attn); free(ane_hidden_attn); free(gate_only_hidden_attn); free(sigmoid_only_hidden_attn); free(v_only_hidden_attn); + free(cpu_attn_q_normed); free(ane_attn_q_normed); free(cpu_attn_gate_sigmoid); free(ane_attn_gate_sigmoid); + free(cpu_attn_k_normed); free(ane_attn_k_normed); free(cpu_attn_q_rope); free(ane_attn_q_rope); + free(cpu_attn_k_rope); free(ane_attn_k_rope); free(cpu_attn_scores); free(ane_attn_scores); + free(cpu_attn_probs); free(ane_attn_probs); free(cpu_attn_context); free(ane_attn_context); + free(cpu_attn_gated_context); free(ane_attn_gated_context); + free(cpu_ffn_rms); free(ane_ffn_rms); free(cpu_ffn_gate); free(ane_ffn_gate); + free(cpu_ffn_up); free(ane_ffn_up); free(ane_ffn_up_only); + free(cpu_ffn_sigmoid); free(cpu_ffn_sigmoid_fp16emu); free(ane_ffn_sigmoid_micro); free(ane_ffn_sigmoid_tanh_micro); + free(cpu_ffn_silu); free(cpu_ffn_silu_fp16emu); free(ane_ffn_silu); free(ane_ffn_silu_only); free(ane_ffn_silu_micro); free(ane_ffn_silu_tanh_micro); + free(cpu_ffn_hidden); free(cpu_ffn_hidden_fp16emu); free(ane_ffn_hidden); free(ane_ffn_hidden_only); free(ane_ffn_mul_micro); + free(cpu_ffn_down); free(ane_ffn_down_same_input); + free(cpu_ffn_final); free(ane_ffn_final_same_input); free(hybrid_layer_final); + if (manifest) orion_qwen35_manifest_free(manifest); + if (tok) orion_gpt2_tokenizer_free(tok); + return 1; + } +} diff --git a/tests/test_qwen35_9b_cpu_infer_smoke.m b/tests/test_qwen35_9b_cpu_infer_smoke.m new file mode 100644 index 0000000..e287902 --- /dev/null +++ b/tests/test_qwen35_9b_cpu_infer_smoke.m @@ -0,0 +1,237 @@ +#import +#import +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +static int file_exists(const char *path) { + return access(path, F_OK) == 0; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static int sampled_topk_logits(const char *blob_dir, + const char *lm_head_name, + const float *hidden, + int d_model, + int sample_vocab, + int *top_id, + float *top_logit) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + int best_id = -1; + float best = -INFINITY; + for (int tok = 0; tok < sample_vocab; tok++) { + if (!orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + if (dot > best) { + best = dot; + best_id = tok; + } + } + + free(row); + *top_id = best_id; + *top_logit = best; + return 1; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int seq_len = 1; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *hidden = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = NULL; + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch || !last_hidden) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(embed_path, 0, d_model, hidden)) { + fprintf(stderr, "FAIL: failed to read token embedding row 0\n"); + goto fail; + } + + for (int layer_idx = 0; layer_idx < manifest->n_layer; layer_idx++) { + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + fprintf(stderr, "FAIL: missing layer %d norm/mlp tensors\n", layer_idx); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + orion_qwen_cpu_rmsnorm(hidden, input_ln, d_model, 1e-6f, normed); + memset(mixer_out, 0, (size_t)d_model * sizeof(float)); + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + if (file_exists(full_q_path)) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + fprintf(stderr, "FAIL: missing full-attention tensors for layer %d\n", layer_idx); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + int qkv_rows = orion_blob_element_count(([NSString stringWithFormat:@"%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx]).UTF8String) / d_model; + int value_dim = orion_blob_element_count(([NSString stringWithFormat:@"%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx]).UTF8String) / d_model; + int num_v_heads = orion_blob_element_count(([NSString stringWithFormat:@"%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx]).UTF8String); + int head_v_dim = orion_blob_element_count(([NSString stringWithFormat:@"%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx]).UTF8String); + int key_dim = (qkv_rows - value_dim) / 2; + int num_k_heads = num_v_heads; + int head_k_dim = key_dim / num_k_heads; + int conv_kernel = orion_blob_element_count(([NSString stringWithFormat:@"%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx]).UTF8String) / qkv_rows; + + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + fprintf(stderr, "FAIL: missing linear-attention tensors for layer %d\n", layer_idx); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < d_model; i++) hidden[i] += mixer_out[i]; + orion_qwen_cpu_rmsnorm(hidden, post_ln, d_model, 1e-6f, scratch); + orion_qwen_cpu_swiglu_ffn(scratch, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out); + for (int i = 0; i < d_model; i++) hidden[i] += mlp_out[i]; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) { + fprintf(stderr, "FAIL: missing final norm\n"); + goto fail; + } + orion_qwen_cpu_rmsnorm(hidden, final_norm, d_model, 1e-6f, last_hidden); + + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + int top_id = -1; + float top_logit = -INFINITY; + if (!sampled_topk_logits(blob_dir, lm_head_name, last_hidden, d_model, 16, &top_id, &top_logit)) { + fprintf(stderr, "FAIL: failed sampled logits scan\n"); + goto fail; + } + + double hidden_abs = abs_sum(hidden, d_model); + double last_hidden_abs = abs_sum(last_hidden, d_model); + if (hidden_abs <= 0.0 || last_hidden_abs <= 0.0 || top_id < 0 || !isfinite(top_logit)) { + fprintf(stderr, "FAIL: expected non-zero hidden/logits path\n"); + goto fail; + } + + printf("PASS: qwen35 9b cpu infer smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" seq_len=%d\n", seq_len); + printf(" layers_run=%d\n", manifest->n_layer); + printf(" hidden_abs_sum=%.6f\n", hidden_abs); + printf(" last_hidden_abs_sum=%.6f\n", last_hidden_abs); + printf(" sampled_top_token_id=%d\n", top_id); + printf(" sampled_top_token_logit=%.6f\n", top_logit); + printf(" next_blocker=%s\n", "9b decode loop and ANE prefill integration"); + + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + free(last_hidden); + free(final_norm); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + free(last_hidden); + free(final_norm); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_9b_decode_loop_cpu_smoke.m b/tests/test_qwen35_9b_decode_loop_cpu_smoke.m new file mode 100644 index 0000000..8c1e6a5 --- /dev/null +++ b/tests/test_qwen35_9b_decode_loop_cpu_smoke.m @@ -0,0 +1,330 @@ +#import +#import +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../tokenizer/gpt2_bpe.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +static NSDictionary* load_json(NSString* path) { + NSData* data = [NSData dataWithContentsOfFile:path]; + if (!data) return nil; + return [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; +} + +static int file_exists(const char *path) { + return access(path, F_OK) == 0; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static int sampled_topk_logits(const char *blob_dir, + const char *lm_head_name, + const float *hidden, + int d_model, + int sample_vocab, + int *top_id, + float *top_logit, + int *runnerup_id, + float *runnerup_logit) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + int best_id = -1; + int second_id = -1; + float best = -INFINITY; + float second = -INFINITY; + for (int tok = 0; tok < sample_vocab; tok++) { + if (!orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + if (dot > best) { + second = best; + second_id = best_id; + best = dot; + best_id = tok; + } else if (dot > second) { + second = dot; + second_id = tok; + } + } + + free(row); + *top_id = best_id; + *top_logit = best; + *runnerup_id = second_id; + *runnerup_logit = second; + return 1; +} + +static int compute_next_token(const char *blob_dir, + OrionQwen35Manifest *manifest, + const int *token_ids, + int seq_len, + int sample_vocab, + int *top_id, + float *top_logit, + int *runnerup_id, + float *runnerup_logit, + double *last_hidden_abs_sum) { + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *hidden = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = NULL; + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch || !last_hidden) goto fail; + + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + for (int s = 0; s < seq_len; s++) { + if (!orion_read_blob_row_f32(embed_path, token_ids[s], d_model, hidden + s * d_model)) goto fail; + } + + for (int layer_idx = 0; layer_idx < manifest->n_layer; layer_idx++) { + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + memset(mixer_out, 0, (size_t)seq_len * d_model * sizeof(float)); + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + if (file_exists(full_q_path)) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + char path_qkv[2048], path_out[2048], path_dt[2048], path_norm[2048], path_conv[2048]; + snprintf(path_qkv, sizeof(path_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + snprintf(path_out, sizeof(path_out), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx); + snprintf(path_dt, sizeof(path_dt), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + snprintf(path_norm, sizeof(path_norm), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx); + snprintf(path_conv, sizeof(path_conv), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + + int qkv_rows = orion_blob_element_count(path_qkv) / d_model; + int value_dim = orion_blob_element_count(path_out) / d_model; + int num_v_heads = orion_blob_element_count(path_dt); + int head_v_dim = orion_blob_element_count(path_norm); + int key_dim = (qkv_rows - value_dim) / 2; + int num_k_heads = num_v_heads; + int head_k_dim = key_dim / num_k_heads; + int conv_kernel = orion_blob_element_count(path_conv) / qkv_rows; + + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) hidden[i] += mixer_out[i]; + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden[i] += mlp_out[i]; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) goto fail; + orion_qwen_cpu_rmsnorm(hidden + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, last_hidden); + + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + if (!sampled_topk_logits(blob_dir, lm_head_name, last_hidden, d_model, sample_vocab, + top_id, top_logit, runnerup_id, runnerup_logit)) goto fail; + *last_hidden_abs_sum = abs_sum(last_hidden, d_model); + + free(hidden); free(normed); free(mixer_out); free(mlp_out); free(scratch); free(last_hidden); free(final_norm); + return 1; + +fail: + free(hidden); free(normed); free(mixer_out); free(mlp_out); free(scratch); free(last_hidden); free(final_norm); + return 0; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + NSString* tokDir = [NSString stringWithUTF8String:argv[2]]; + NSDictionary* meta = load_json([tokDir stringByAppendingPathComponent:@"meta.json"]); + if (!meta) { + fprintf(stderr, "FAIL: missing tokenizer meta.json\n"); + return 1; + } + NSString* regex = meta[@"regex_pattern"]; + NSString* vocabPath = [tokDir stringByAppendingPathComponent:@"vocab.json"]; + NSString* mergesPath = [tokDir stringByAppendingPathComponent:@"merges.txt"]; + OrionGPT2Tokenizer* tok = orion_gpt2_tokenizer_load_with_regex(vocabPath.UTF8String, mergesPath.UTF8String, regex.UTF8String); + if (!tok) { + fprintf(stderr, "FAIL: tokenizer load failed\n"); + return 1; + } + + OrionQwen35Manifest* manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const char *prompt = getenv("PROMPT"); + if (!prompt || !prompt[0]) prompt = "YES"; + int tokens[64] = {0}; + int prompt_len = orion_gpt2_encode(tok, prompt, tokens, 64); + if (prompt_len <= 0) { + fprintf(stderr, "FAIL: prompt encode failed\n"); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const int sample_vocab = 4096; + int gen_ids[2] = {0}; + float gen_logits[2] = {0.0f, 0.0f}; + int gen_runnerup_ids[2] = {0}; + float gen_runnerup_logits[2] = {0.0f, 0.0f}; + float gen_top2_margins[2] = {0.0f, 0.0f}; + double final_hidden_abs = 0.0; + int seq[128] = {0}; + memcpy(seq, tokens, (size_t)prompt_len * sizeof(int)); + int seq_len = prompt_len; + + for (int step = 0; step < 2; step++) { + int next_id = -1; + float next_logit = -INFINITY; + int runnerup_id = -1; + float runnerup_logit = -INFINITY; + if (!compute_next_token(blob_dir, manifest, seq, seq_len, sample_vocab, + &next_id, &next_logit, &runnerup_id, &runnerup_logit, &final_hidden_abs)) { + fprintf(stderr, "FAIL: top1 compute failed at step %d\n", step); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + gen_ids[step] = next_id; + gen_logits[step] = next_logit; + gen_runnerup_ids[step] = runnerup_id; + gen_runnerup_logits[step] = runnerup_logit; + gen_top2_margins[step] = next_logit - runnerup_logit; + seq[seq_len++] = next_id; + } + + char* decoded_prompt = orion_gpt2_decode(tok, tokens, prompt_len); + char* decoded_generated = orion_gpt2_decode(tok, gen_ids, 2); + char* decoded_full = orion_gpt2_decode(tok, seq, seq_len); + if (!decoded_prompt || !decoded_generated || !decoded_full) { + fprintf(stderr, "FAIL: decode returned NULL\n"); + free(decoded_prompt); free(decoded_generated); free(decoded_full); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + printf("PASS: qwen35 9b decode loop cpu smoke\n"); + printf(" prompt=%s\n", prompt); + printf(" prompt_len=%d\n", prompt_len); + printf(" gen_len=%d\n", 2); + printf(" sample_vocab=%d\n", sample_vocab); + printf(" gen_token_0=%d\n", gen_ids[0]); + printf(" gen_token_1=%d\n", gen_ids[1]); + printf(" gen_logit_0=%.6f\n", gen_logits[0]); + printf(" gen_logit_1=%.6f\n", gen_logits[1]); + printf(" gen_runnerup_token_0=%d\n", gen_runnerup_ids[0]); + printf(" gen_runnerup_token_1=%d\n", gen_runnerup_ids[1]); + printf(" gen_runnerup_logit_0=%.6f\n", gen_runnerup_logits[0]); + printf(" gen_runnerup_logit_1=%.6f\n", gen_runnerup_logits[1]); + printf(" gen_top2_margin_0=%.6f\n", gen_top2_margins[0]); + printf(" gen_top2_margin_1=%.6f\n", gen_top2_margins[1]); + printf(" final_hidden_abs_sum=%.6f\n", final_hidden_abs); + printf(" decoded_prompt=%s\n", decoded_prompt); + printf(" decoded_generated=%s\n", decoded_generated); + printf(" decoded_full=%s\n", decoded_full); + printf(" next_blocker=%s\n", "9b ANE prefill integration"); + + free(decoded_prompt); + free(decoded_generated); + free(decoded_full); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 0; + } +} diff --git a/tests/test_qwen35_9b_decode_loop_mixed_ane_cpu_smoke.m b/tests/test_qwen35_9b_decode_loop_mixed_ane_cpu_smoke.m new file mode 100644 index 0000000..5b58472 --- /dev/null +++ b/tests/test_qwen35_9b_decode_loop_mixed_ane_cpu_smoke.m @@ -0,0 +1,1201 @@ +#import +#import +#include +#include +#include +#include +#include +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#include "core/bucket.h" +#import "../model/weight_loader.h" +#import "../tokenizer/gpt2_bpe.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static const int kQwenPrefillBuckets[] = {32, 64, 128, 256, 512, 1024}; +static const int kQwenPrefillNumBuckets = (int)(sizeof(kQwenPrefillBuckets) / sizeof(kQwenPrefillBuckets[0])); + +static NSDictionary* load_json(NSString* path) { + NSData* data = [NSData dataWithContentsOfFile:path]; + if (!data) return nil; + return [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; +} + +static int file_exists(const char *path) { + return access(path, F_OK) == 0; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +static void cpu_linear_batch(const float *x_seq, + int seq_len, + const float *weight, + int in_dim, + int out_dim, + float *out_seq) { + for (int s = 0; s < seq_len; s++) { + const float *x = x_seq + (size_t)s * in_dim; + float *out = out_seq + (size_t)s * out_dim; + for (int o = 0; o < out_dim; o++) { + const float *w = weight + (size_t)o * in_dim; + float acc = 0.0f; + for (int i = 0; i < in_dim; i++) acc += x[i] * w[i]; + out[o] = acc; + } + } +} + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_qproj_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_down_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_down_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +typedef struct { + OrionProgram *prog_q; + OrionProgram *prog_kv; + OrionProgram *prog_ffn; + int layer_idx; + int bucket; + int d_model; + int q_dim; + int kv_dim; + int qkv_input_mode; + int q_uses_cpu_rms; + int kv_uses_cpu_rms; + int q_proj_uses_cpu; + int k_proj_uses_cpu; + int v_proj_uses_cpu; + int q_query_uses_cpu; + int q_gate_uses_cpu; + int ffn_uses_cpu; + int q_gate_cpu_channel_count; + int q_gate_cpu_channels[32]; + int v_proj_cpu_channel_count; + int v_proj_cpu_channels[32]; +} OrionQwen35AneBridge; + +enum { + ORION_QKV_INPUT_MODE_ANE_RMS = 0, + ORION_QKV_INPUT_MODE_CPU_RMS = 1, + ORION_QKV_INPUT_MODE_KV_CPU_RMS = 2, +}; + +static int parse_index_list_local(const char *csv, int *out, int max_count, int limit) { + if (!csv || !*csv || !out || max_count <= 0) return 0; + char *copy = strdup(csv); + if (!copy) return 0; + int count = 0; + char *save = NULL; + for (char *tok = strtok_r(copy, ",", &save); tok && count < max_count; tok = strtok_r(NULL, ",", &save)) { + while (*tok && isspace((unsigned char)*tok)) tok++; + if (!*tok) continue; + char *end = NULL; + long value = strtol(tok, &end, 10); + if (end == tok || value < 0 || value >= limit) continue; + int dup = 0; + for (int i = 0; i < count; i++) { + if (out[i] == (int)value) { + dup = 1; + break; + } + } + if (!dup) out[count++] = (int)value; + } + free(copy); + return count; +} + +static int override_enabled_for_layer(const char *env_name, int layer, int layer_limit) { + const char *csv = getenv(env_name); + if (!csv || !*csv) return 1; + int layers[64] = {0}; + int count = parse_index_list_local(csv, layers, (int)(sizeof(layers) / sizeof(layers[0])), layer_limit); + if (count <= 0) return 1; + for (int i = 0; i < count; i++) { + if (layers[i] == layer) return 1; + } + return 0; +} + +static int load_qkv_input_mode(void) { + const char *mode = getenv("ORION_QKV_INPUT_MODE"); + if (!mode || !*mode) return ORION_QKV_INPUT_MODE_ANE_RMS; + if (strcmp(mode, "cpu_rms") == 0) return ORION_QKV_INPUT_MODE_CPU_RMS; + if (strcmp(mode, "kv_cpu_rms") == 0) return ORION_QKV_INPUT_MODE_KV_CPU_RMS; + return ORION_QKV_INPUT_MODE_ANE_RMS; +} + +static int q_uses_cpu_rms_mode(int qkv_input_mode) { + return qkv_input_mode == ORION_QKV_INPUT_MODE_CPU_RMS; +} + +static int kv_uses_cpu_rms_mode(int qkv_input_mode) { + return qkv_input_mode == ORION_QKV_INPUT_MODE_CPU_RMS || + qkv_input_mode == ORION_QKV_INPUT_MODE_KV_CPU_RMS; +} + +static const char *qkv_input_mode_label(int qkv_input_mode) { + switch (qkv_input_mode) { + case ORION_QKV_INPUT_MODE_CPU_RMS: + return "cpu_rms_linear_only"; + case ORION_QKV_INPUT_MODE_KV_CPU_RMS: + return "kv_cpu_rms_linear_only"; + default: + return "ane_rms_plus_linear"; + } +} + +static int use_cpu_v_proj_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_k_proj_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_K_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_matmul_v_proj_mode(void) { + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "matmul") == 0; +} + +static int load_v_proj_cpu_channel_override(int *out_channels, int max_channels, int q_dim) { + const char *csv = getenv("ORION_V_PROJ_CPU_CHANNELS"); + if (!csv || !*csv) return 0; + return parse_index_list_local(csv, out_channels, max_channels, q_dim); +} + +static int use_cpu_q_proj_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_Q_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_query_override(void) { + const char *source = getenv("ORION_Q_QUERY_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_gate_override(void) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_ffn_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_FFN_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int load_q_gate_cpu_channel_preset(const char *source, int *out_channels, int max_channels, int q_dim) { + static const int kSeed2[] = {143, 3994}; + static const int kSeed4[] = {143, 3994, 3768, 1353}; + const int *preset = NULL; + int preset_count = 0; + if (!source || !out_channels || max_channels <= 0) return 0; + if (strcmp(source, "seed2") == 0) { + preset = kSeed2; + preset_count = (int)(sizeof(kSeed2) / sizeof(kSeed2[0])); + } else if (strcmp(source, "seed4") == 0) { + preset = kSeed4; + preset_count = (int)(sizeof(kSeed4) / sizeof(kSeed4[0])); + } + if (!preset) return 0; + int count = 0; + for (int i = 0; i < preset_count && count < max_channels; i++) { + int channel = preset[i]; + if (channel < 0 || channel >= q_dim) continue; + out_channels[count++] = channel; + } + return count; +} + +static int load_q_gate_cpu_channel_override(int *out_channels, int max_channels, int q_dim) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + int preset_count = load_q_gate_cpu_channel_preset(source, out_channels, max_channels, q_dim); + if (preset_count > 0) return preset_count; + const char *csv = getenv("ORION_Q_GATE_CPU_CHANNELS"); + if (!csv || !*csv) return 0; + return parse_index_list_local(csv, out_channels, max_channels, q_dim); +} + +static const char *q_proj_source_label(int q_proj_uses_cpu) { + return q_proj_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *q_query_source_label(int q_query_uses_cpu) { + return q_query_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *k_proj_source_label(int k_proj_uses_cpu) { + return k_proj_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *q_gate_source_label(const char *source, int q_gate_uses_cpu, int partial_channel_count) { + if (q_gate_uses_cpu) return "cpu_linear"; + if (source && strcmp(source, "seed2") == 0) return "ane_linear_seed2"; + if (source && strcmp(source, "seed4") == 0) return "ane_linear_seed4"; + if (partial_channel_count > 0) return "ane_linear_cpu_channels"; + return "ane_linear"; +} + +static const char *ffn_source_label(int ffn_uses_cpu) { + return ffn_uses_cpu ? "cpu" : "ane"; +} + +static const char *v_proj_source_label(int v_proj_uses_cpu, int partial_channel_count) { + if (v_proj_uses_cpu) return "cpu_linear"; + if (partial_channel_count > 0) return "ane_linear_cpu_channels"; + if (use_matmul_v_proj_mode()) return "ane_matmul"; + return "ane_linear"; +} + +static void apply_v_proj_cpu_channel_overrides(float *dst_v_proj_seq, + const float *cpu_v_proj_seq, + int seq_len, + int kv_dim, + int head_dim, + int q_per_kv, + const int *channels, + int channel_count) { + if (!dst_v_proj_seq || !cpu_v_proj_seq || !channels || channel_count <= 0 || seq_len <= 0 || + kv_dim <= 0 || head_dim <= 0 || q_per_kv <= 0) { + return; + } + for (int i = 0; i < channel_count; i++) { + int channel = channels[i]; + int head = channel / head_dim; + int offset = channel % head_dim; + int kv_head = head / q_per_kv; + int kv_channel = kv_head * head_dim + offset; + if (kv_channel < 0 || kv_channel >= kv_dim) continue; + for (int s = 0; s < seq_len; s++) { + dst_v_proj_seq[(size_t)s * kv_dim + kv_channel] = + cpu_v_proj_seq[(size_t)s * kv_dim + kv_channel]; + } + } +} + +static void apply_q_gate_cpu_channel_overrides(float *dst_q_proj_seq, + const float *cpu_q_proj_seq, + int seq_len, + int q_dim, + const int *channels, + int channel_count) { + if (!dst_q_proj_seq || !cpu_q_proj_seq || !channels || channel_count <= 0 || seq_len <= 0 || q_dim <= 0) { + return; + } + for (int i = 0; i < channel_count; i++) { + int channel = channels[i]; + if (channel < 0 || channel >= q_dim) continue; + for (int s = 0; s < seq_len; s++) { + dst_q_proj_seq[(size_t)s * (q_dim * 2) + q_dim + channel] = + cpu_q_proj_seq[(size_t)s * (q_dim * 2) + q_dim + channel]; + } + } +} + +static int load_full_attention_layers(OrionQwen35Manifest *manifest, int *out_layers, int max_layers) { + if (!manifest || !manifest->manifest_path || !out_layers || max_layers <= 0) return 0; + NSString *manifestPath = [NSString stringWithUTF8String:manifest->manifest_path]; + NSDictionary *root = load_json(manifestPath); + NSDictionary *runtime = [root isKindOfClass:[NSDictionary class]] ? root[@"runtime"] : nil; + NSArray *layerTypes = [runtime isKindOfClass:[NSDictionary class]] ? runtime[@"layer_types"] : nil; + if (![layerTypes isKindOfClass:[NSArray class]]) return 0; + + int count = 0; + for (NSInteger i = 0; i < [layerTypes count] && count < max_layers; i++) { + id value = layerTypes[i]; + if ([value isKindOfClass:[NSString class]] && [(NSString *)value isEqualToString:@"full_attention"]) { + out_layers[count++] = (int)i; + } + } + return count; +} + +static void bridge_release(OrionQwen35AneBridge *bridge) { + if (!bridge) return; + if (bridge->prog_q) orion_release_program(bridge->prog_q); + if (bridge->prog_kv) orion_release_program(bridge->prog_kv); + if (bridge->prog_ffn) orion_release_program(bridge->prog_ffn); + memset(bridge, 0, sizeof(*bridge)); +} + +static int bridge_init(OrionQwen35AneBridge *bridge, + NSString *blobDir, + int layer, + int bucket, + OrionQwen35Manifest *manifest, + int qkv_input_mode, + int q_proj_uses_cpu, + int k_proj_uses_cpu, + int v_proj_uses_cpu, + int q_query_uses_cpu, + int q_gate_uses_cpu, + int ffn_uses_cpu) { + memset(bridge, 0, sizeof(*bridge)); + const int q_uses_cpu_rms = q_uses_cpu_rms_mode(qkv_input_mode); + const int kv_uses_cpu_rms = kv_uses_cpu_rms_mode(qkv_input_mode); + OrionModelConfig cfg = { + .n_layer = manifest->n_layer, + .n_head = manifest->n_head, + .n_kv_head = manifest->n_kv_head, + .d_model = manifest->d_model, + .head_dim = manifest->head_dim, + .hidden_dim = manifest->d_ff, + .vocab = manifest->vocab, + .max_seq = manifest->max_seq, + }; + NSString *mil_q = compile_graph( + q_uses_cpu_rms + ? orion_frontend_qwen35_prefill_q_proj_linear_only(layer, bucket, &cfg) + : orion_frontend_qwen35_prefill_q_proj(layer, bucket, &cfg) + ); + NSString *mil_kv = compile_graph( + kv_uses_cpu_rms + ? orion_frontend_qwen35_prefill_kv_proj_linear_only(layer, bucket, &cfg) + : orion_frontend_qwen35_prefill_kv_proj(layer, bucket, &cfg) + ); + NSString *mil_ffn = compile_graph(orion_frontend_qwen35_prefill_ffn(layer, bucket, &cfg)); + if (!mil_q || !mil_kv || !mil_ffn) return 0; + + bridge->prog_q = orion_compile_mil( + mil_q.UTF8String, + q_uses_cpu_rms ? build_qproj_linear_only_wdict(layer, blobDir) : build_qproj_wdict(layer, blobDir), + "qwen35_9b_decode_q" + ); + bridge->prog_kv = orion_compile_mil( + mil_kv.UTF8String, + kv_uses_cpu_rms ? build_kv_linear_only_wdict(layer, blobDir) : build_kv_wdict(layer, blobDir), + "qwen35_9b_decode_kv" + ); + bridge->prog_ffn = orion_compile_mil(mil_ffn.UTF8String, build_ffn_wdict(layer, blobDir), "qwen35_9b_decode_ffn"); + if (!bridge->prog_q || !bridge->prog_kv || !bridge->prog_ffn) { + bridge_release(bridge); + return 0; + } + + bridge->layer_idx = layer; + bridge->bucket = bucket; + bridge->d_model = manifest->d_model; + bridge->q_dim = manifest->n_head * manifest->head_dim; + bridge->kv_dim = manifest->n_kv_head * manifest->head_dim; + bridge->qkv_input_mode = qkv_input_mode; + bridge->q_uses_cpu_rms = q_uses_cpu_rms; + bridge->kv_uses_cpu_rms = kv_uses_cpu_rms; + bridge->q_proj_uses_cpu = q_proj_uses_cpu; + bridge->k_proj_uses_cpu = k_proj_uses_cpu; + bridge->v_proj_uses_cpu = v_proj_uses_cpu; + bridge->q_query_uses_cpu = q_query_uses_cpu; + bridge->q_gate_uses_cpu = q_gate_uses_cpu; + bridge->ffn_uses_cpu = ffn_uses_cpu; + bridge->q_gate_cpu_channel_count = load_q_gate_cpu_channel_override( + bridge->q_gate_cpu_channels, + (int)(sizeof(bridge->q_gate_cpu_channels) / sizeof(bridge->q_gate_cpu_channels[0])), + bridge->q_dim + ); + if (bridge->q_gate_cpu_channel_count > 0 && + !override_enabled_for_layer("ORION_Q_GATE_CPU_CHANNEL_LAYERS", layer, manifest->n_layer)) { + bridge->q_gate_cpu_channel_count = 0; + } + if (bridge->ffn_uses_cpu && + !override_enabled_for_layer("ORION_FFN_CPU_LAYERS", layer, manifest->n_layer)) { + bridge->ffn_uses_cpu = 0; + } + bridge->v_proj_cpu_channel_count = load_v_proj_cpu_channel_override( + bridge->v_proj_cpu_channels, + (int)(sizeof(bridge->v_proj_cpu_channels) / sizeof(bridge->v_proj_cpu_channels[0])), + bridge->q_dim + ); + return 1; +} + +static int mixed_full_attention_layer(const char *blob_dir, + OrionQwen35Manifest *manifest, + OrionQwen35AneBridge *bridge, + const float *hidden_in, + int seq_len, + float *hidden_out) { + const int layer = bridge->layer_idx; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = bridge->q_dim; + const int kv_dim = bridge->kv_dim; + const int full_cpu_proj_path = + bridge->q_proj_uses_cpu && + bridge->k_proj_uses_cpu && + bridge->v_proj_uses_cpu && + bridge->ffn_uses_cpu; + + float *input_ln = NULL; + float *post_ln = NULL; + float *gate_proj = NULL; + float *up_proj = NULL; + float *down_proj = NULL; + float *o_proj = NULL; + float *q_norm = NULL; + float *k_norm = NULL; + float *q_proj = NULL; + float *k_proj = NULL; + float *v_proj = NULL; + float *normed = NULL; + float *cpu_k_proj_seq = NULL; + float *cpu_v_proj_seq = NULL; + float *q_proj_seq = NULL; + float *k_proj_seq = NULL; + float *v_proj_seq = NULL; + float *attn = NULL; + float *hidden_attn = NULL; + float *post_normed = NULL; + float *mlp = NULL; + IOSurfaceRef ioInQ = NULL; + IOSurfaceRef ioInKV = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + IOSurfaceRef ioFfnIn = NULL; + IOSurfaceRef ioHidden = NULL; + + input_ln = load_exact(blob_dir, layer, "input_layernorm.bin", d_model); + post_ln = load_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + gate_proj = load_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + up_proj = load_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + down_proj = load_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + o_proj = load_exact(blob_dir, layer, "self_attn_o_proj.bin", d_model * q_dim); + q_norm = load_exact(blob_dir, layer, "self_attn_q_norm.bin", head_dim); + k_norm = load_exact(blob_dir, layer, "self_attn_k_norm.bin", head_dim); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0 || full_cpu_proj_path) { + q_proj = load_exact(blob_dir, layer, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + } + if (bridge->k_proj_uses_cpu || full_cpu_proj_path) { + k_proj = load_exact(blob_dir, layer, "self_attn_k_proj.bin", kv_dim * d_model); + } + if (bridge->v_proj_uses_cpu || bridge->v_proj_cpu_channel_count > 0 || full_cpu_proj_path) { + v_proj = load_exact(blob_dir, layer, "self_attn_v_proj.bin", kv_dim * d_model); + } + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj || !o_proj || !q_norm || !k_norm || + ((bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0 || full_cpu_proj_path) && !q_proj) || + ((bridge->k_proj_uses_cpu || full_cpu_proj_path) && !k_proj) || + ((bridge->v_proj_uses_cpu || bridge->v_proj_cpu_channel_count > 0 || full_cpu_proj_path) && !v_proj)) goto fail; + + normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + cpu_k_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + q_proj_seq = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + k_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + v_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + post_normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + mlp = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!normed || !cpu_k_proj_seq || !cpu_v_proj_seq || !q_proj_seq || !k_proj_seq || !v_proj_seq || !attn || !hidden_attn || !post_normed || !mlp) goto fail; + + for (int t = 0; t < seq_len; t++) { + orion_qwen_cpu_rmsnorm(hidden_in + t * d_model, input_ln, d_model, 1e-6f, normed + t * d_model); + } + + if (full_cpu_proj_path) { + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn + ); + + for (int i = 0; i < seq_len * d_model; i++) hidden_attn[i] = hidden_in[i] + attn[i]; + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn + s * d_model, post_ln, d_model, 1e-6f, post_normed + s * d_model); + orion_qwen_cpu_swiglu_ffn(post_normed + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_out[i] = hidden_attn[i] + mlp[i]; + + if (input_ln) free(input_ln); + if (post_ln) free(post_ln); + if (gate_proj) free(gate_proj); + if (up_proj) free(up_proj); + if (down_proj) free(down_proj); + if (o_proj) free(o_proj); + if (q_norm) free(q_norm); + if (k_norm) free(k_norm); + if (q_proj) free(q_proj); + if (k_proj) free(k_proj); + if (v_proj) free(v_proj); + if (normed) free(normed); + if (cpu_k_proj_seq) free(cpu_k_proj_seq); + if (cpu_v_proj_seq) free(cpu_v_proj_seq); + if (q_proj_seq) free(q_proj_seq); + if (k_proj_seq) free(k_proj_seq); + if (v_proj_seq) free(v_proj_seq); + if (attn) free(attn); + if (hidden_attn) free(hidden_attn); + if (post_normed) free(post_normed); + if (mlp) free(mlp); + return 1; + } + + const float *q_input = bridge->q_uses_cpu_rms ? normed : hidden_in; + const float *kv_input = bridge->kv_uses_cpu_rms ? normed : hidden_in; + ioInQ = make_cpu_seq_input_surface(q_input, seq_len, bridge->bucket, d_model); + ioInKV = make_cpu_seq_input_surface(kv_input, seq_len, bridge->bucket, d_model); + ioQ = make_f32_surface((q_dim * 2) * bridge->bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bridge->bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bridge->bucket, 0.0f); + + IOSurfaceRef insQ[] = {ioInQ}; + IOSurfaceRef outsQ[] = {ioQ}; + IOSurfaceRef insKV[] = {ioInKV}; + IOSurfaceRef outsKV[] = {ioK, ioV}; + if (!orion_eval(bridge->prog_q, insQ, 1, outsQ, 1) || !orion_eval(bridge->prog_kv, insKV, 1, outsKV, 2)) goto fail; + + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bridge->bucket, q_proj_seq); + read_ane_surface_prefix(ioK, kv_dim, seq_len, bridge->bucket, k_proj_seq); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bridge->bucket, v_proj_seq); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0) { + float *cpu_q_proj_seq = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + if (!cpu_q_proj_seq) goto fail; + cpu_linear_batch(normed, seq_len, q_proj, d_model, q_dim * 2, cpu_q_proj_seq); + for (int s = 0; s < seq_len; s++) { + float *dst = q_proj_seq + (size_t)s * (q_dim * 2); + const float *src = cpu_q_proj_seq + (size_t)s * (q_dim * 2); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu) { + memcpy(dst, src, (size_t)q_dim * sizeof(float)); + } + if (bridge->q_proj_uses_cpu || bridge->q_gate_uses_cpu) { + memcpy(dst + q_dim, src + q_dim, (size_t)q_dim * sizeof(float)); + } + } + if (!bridge->q_proj_uses_cpu && !bridge->q_gate_uses_cpu && bridge->q_gate_cpu_channel_count > 0) { + apply_q_gate_cpu_channel_overrides(q_proj_seq, cpu_q_proj_seq, seq_len, q_dim, + bridge->q_gate_cpu_channels, bridge->q_gate_cpu_channel_count); + } + free(cpu_q_proj_seq); + } + if (bridge->k_proj_uses_cpu) { + cpu_linear_batch(normed, seq_len, k_proj, d_model, kv_dim, cpu_k_proj_seq); + memcpy(k_proj_seq, cpu_k_proj_seq, (size_t)seq_len * kv_dim * sizeof(float)); + } + if (bridge->v_proj_uses_cpu || bridge->v_proj_cpu_channel_count > 0) { + cpu_linear_batch(normed, seq_len, v_proj, d_model, kv_dim, cpu_v_proj_seq); + } + if (bridge->v_proj_uses_cpu) { + memcpy(v_proj_seq, cpu_v_proj_seq, (size_t)seq_len * kv_dim * sizeof(float)); + } else if (bridge->v_proj_cpu_channel_count > 0) { + const int q_per_kv = n_head / n_kv_head; + apply_v_proj_cpu_channel_overrides(v_proj_seq, cpu_v_proj_seq, seq_len, kv_dim, head_dim, q_per_kv, + bridge->v_proj_cpu_channels, bridge->v_proj_cpu_channel_count); + } + + orion_qwen_cpu_full_attention_from_projections_with_rope( + q_proj_seq, k_proj_seq, v_proj_seq, seq_len, + o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + attn + ); + + for (int i = 0; i < seq_len * d_model; i++) hidden_attn[i] = hidden_in[i] + attn[i]; + if (bridge->ffn_uses_cpu) { + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn + s * d_model, post_ln, d_model, 1e-6f, post_normed + s * d_model); + orion_qwen_cpu_swiglu_ffn(post_normed + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_out[i] = hidden_attn[i] + mlp[i]; + } else { + ioFfnIn = make_cpu_seq_input_surface(hidden_attn, seq_len, bridge->bucket, d_model); + ioHidden = make_f32_surface(d_model * bridge->bucket, 0.0f); + IOSurfaceRef insFFN[] = {ioFfnIn}; + IOSurfaceRef outsFFN[] = {ioHidden}; + if (!orion_eval(bridge->prog_ffn, insFFN, 1, outsFFN, 1)) goto fail; + read_ane_surface_prefix(ioHidden, d_model, seq_len, bridge->bucket, hidden_out); + } + + if (ioInQ) CFRelease(ioInQ); + if (ioInKV) CFRelease(ioInKV); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioHidden) CFRelease(ioHidden); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); free(o_proj); free(q_norm); free(k_norm); free(q_proj); free(k_proj); free(v_proj); + free(normed); free(cpu_k_proj_seq); free(cpu_v_proj_seq); free(q_proj_seq); free(k_proj_seq); free(v_proj_seq); free(attn); free(hidden_attn); free(post_normed); free(mlp); + return 1; + +fail: + if (ioInQ) CFRelease(ioInQ); + if (ioInKV) CFRelease(ioInKV); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioHidden) CFRelease(ioHidden); + if (input_ln) free(input_ln); + if (post_ln) free(post_ln); + if (gate_proj) free(gate_proj); + if (up_proj) free(up_proj); + if (down_proj) free(down_proj); + if (o_proj) free(o_proj); + if (q_norm) free(q_norm); + if (k_norm) free(k_norm); + if (q_proj) free(q_proj); + if (k_proj) free(k_proj); + if (v_proj) free(v_proj); + if (normed) free(normed); + if (cpu_k_proj_seq) free(cpu_k_proj_seq); + if (cpu_v_proj_seq) free(cpu_v_proj_seq); + if (q_proj_seq) free(q_proj_seq); + if (k_proj_seq) free(k_proj_seq); + if (v_proj_seq) free(v_proj_seq); + if (attn) free(attn); + if (hidden_attn) free(hidden_attn); + if (post_normed) free(post_normed); + if (mlp) free(mlp); + return 0; +} + +static int sampled_topk_logits(const char *blob_dir, + const char *lm_head_name, + const float *hidden, + int d_model, + int sample_vocab, + int *top_id, + float *top_logit, + int *runnerup_id, + float *runnerup_logit) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + int best_id = -1; + int second_id = -1; + float best = -INFINITY; + float second = -INFINITY; + for (int tok = 0; tok < sample_vocab; tok++) { + if (!orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + if (dot > best) { + second = best; + second_id = best_id; + best = dot; + best_id = tok; + } else if (dot > second) { + second = dot; + second_id = tok; + } + } + + free(row); + *top_id = best_id; + *top_logit = best; + *runnerup_id = second_id; + *runnerup_logit = second; + return 1; +} + +static int compute_next_token_mixed(const char *blob_dir, + OrionQwen35Manifest *manifest, + OrionQwen35AneBridge *bridges, + const unsigned char *bridge_mask, + const int *token_ids, + int seq_len, + int sample_vocab, + const char *lm_head_name_override, + int *top_id, + float *top_logit, + int *runnerup_id, + float *runnerup_logit, + double *last_hidden_abs_sum) { + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + + float *hidden = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = NULL; + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch || !last_hidden) goto fail; + + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + for (int s = 0; s < seq_len; s++) { + if (!orion_read_blob_row_f32(embed_path, token_ids[s], d_model, hidden + s * d_model)) goto fail; + } + + for (int layer_idx = 0; layer_idx < manifest->n_layer; layer_idx++) { + if (bridge_mask && bridge_mask[layer_idx]) { + if (!mixed_full_attention_layer(blob_dir, manifest, &bridges[layer_idx], hidden, seq_len, mixer_out)) goto fail; + memcpy(hidden, mixer_out, (size_t)seq_len * d_model * sizeof(float)); + continue; + } + + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + memset(mixer_out, 0, (size_t)seq_len * d_model * sizeof(float)); + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + if (file_exists(full_q_path)) { + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + char path_qkv[2048], path_out[2048], path_dt[2048], path_norm[2048], path_conv[2048]; + snprintf(path_qkv, sizeof(path_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + snprintf(path_out, sizeof(path_out), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx); + snprintf(path_dt, sizeof(path_dt), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + snprintf(path_norm, sizeof(path_norm), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx); + snprintf(path_conv, sizeof(path_conv), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + + int qkv_rows = orion_blob_element_count(path_qkv) / d_model; + int value_dim = orion_blob_element_count(path_out) / d_model; + int num_v_heads = orion_blob_element_count(path_dt); + int head_v_dim = orion_blob_element_count(path_norm); + int key_dim = (qkv_rows - value_dim) / 2; + int num_k_heads = num_v_heads; + int head_k_dim = key_dim / num_k_heads; + int conv_kernel = orion_blob_element_count(path_conv) / qkv_rows; + + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) hidden[i] += mixer_out[i]; + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden[i] += mlp_out[i]; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) goto fail; + orion_qwen_cpu_rmsnorm(hidden + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, last_hidden); + + const char *lm_head_name = lm_head_name_override + ? lm_head_name_override + : (manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"); + if (!sampled_topk_logits(blob_dir, lm_head_name, last_hidden, d_model, sample_vocab, + top_id, top_logit, runnerup_id, runnerup_logit)) goto fail; + *last_hidden_abs_sum = abs_sum(last_hidden, d_model); + + free(hidden); free(normed); free(mixer_out); free(mlp_out); free(scratch); free(last_hidden); free(final_norm); + return 1; + +fail: + free(hidden); free(normed); free(mixer_out); free(mlp_out); free(scratch); free(last_hidden); free(final_norm); + return 0; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s [single|all_full] [prompt] [gen_len] [sample_vocab] [lm_head_name]\n", argv[0]); + return 2; + } + if (!orion_ane_init()) { + fprintf(stderr, "FAIL: orion_ane_init failed\n"); + return 3; + } + + const char *blob_dir = argv[1]; + NSString* tokDir = [NSString stringWithUTF8String:argv[2]]; + NSDictionary* meta = load_json([tokDir stringByAppendingPathComponent:@"meta.json"]); + if (!meta) { + fprintf(stderr, "FAIL: missing tokenizer meta.json\n"); + return 1; + } + NSString* regex = meta[@"regex_pattern"]; + NSString* vocabPath = [tokDir stringByAppendingPathComponent:@"vocab.json"]; + NSString* mergesPath = [tokDir stringByAppendingPathComponent:@"merges.txt"]; + OrionGPT2Tokenizer* tok = orion_gpt2_tokenizer_load_with_regex(vocabPath.UTF8String, mergesPath.UTF8String, regex.UTF8String); + if (!tok) { + fprintf(stderr, "FAIL: tokenizer load failed\n"); + return 1; + } + + OrionQwen35Manifest* manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const char *mode = (argc >= 4) ? argv[3] : "single"; + const int qkv_input_mode = load_qkv_input_mode(); + const int q_proj_uses_cpu = use_cpu_q_proj_override(); + const int k_proj_uses_cpu = use_cpu_k_proj_override(); + const int q_query_uses_cpu = use_cpu_q_query_override(); + const int q_gate_uses_cpu = use_cpu_q_gate_override(); + const int ffn_uses_cpu = use_cpu_ffn_override(); + const char *q_gate_source = getenv("ORION_Q_GATE_SOURCE"); + const int v_proj_uses_cpu = use_cpu_v_proj_override(); + int q_gate_cpu_channels[32] = {0}; + const int q_gate_cpu_channel_count = load_q_gate_cpu_channel_override( + q_gate_cpu_channels, + (int)(sizeof(q_gate_cpu_channels) / sizeof(q_gate_cpu_channels[0])), + manifest->n_head * manifest->head_dim + ); + int v_proj_cpu_channels[32] = {0}; + const int v_proj_cpu_channel_count = load_v_proj_cpu_channel_override( + v_proj_cpu_channels, + (int)(sizeof(v_proj_cpu_channels) / sizeof(v_proj_cpu_channels[0])), + manifest->n_head * manifest->head_dim + ); + + const char *prompt = (argc >= 5) ? argv[4] : "YES"; + int gen_len = (argc >= 6) ? atoi(argv[5]) : 2; + int sample_vocab = (argc >= 7) ? atoi(argv[6]) : 4096; + const char *lm_head_name_override = (argc >= 8) ? argv[7] : NULL; + if (gen_len <= 0 || gen_len > 8) gen_len = 2; + if (sample_vocab <= 0) sample_vocab = 4096; + int tokens[64] = {0}; + int prompt_len = orion_gpt2_encode(tok, prompt, tokens, 64); + if (prompt_len <= 0) { + fprintf(stderr, "FAIL: prompt encode failed\n"); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const int required_seq_len = prompt_len + gen_len - 1; + const int bucket = orion_select_bucket(required_seq_len, kQwenPrefillBuckets, kQwenPrefillNumBuckets); + if (bucket <= 0) { + fprintf(stderr, "FAIL: no bucket fits required_seq_len=%d\n", required_seq_len); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + OrionQwen35AneBridge *bridges = calloc((size_t)manifest->n_layer, sizeof(OrionQwen35AneBridge)); + unsigned char *bridge_mask = calloc((size_t)manifest->n_layer, sizeof(unsigned char)); + int bridged_layer = 3; + int bridged_layer_count = 0; + if (!bridges || !bridge_mask) { + fprintf(stderr, "FAIL: bridge allocation failed\n"); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + if (strcmp(mode, "all_full") == 0) { + int full_layers[128] = {0}; + int full_count = load_full_attention_layers(manifest, full_layers, 128); + if (full_count <= 0) { + fprintf(stderr, "FAIL: no full_attention layers found\n"); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + bridged_layer = full_layers[0]; + for (int i = 0; i < full_count; i++) { + int layer = full_layers[i]; + if (!bridge_init(&bridges[layer], [NSString stringWithUTF8String:blob_dir], layer, bucket, manifest, qkv_input_mode, q_proj_uses_cpu, k_proj_uses_cpu, v_proj_uses_cpu, q_query_uses_cpu, q_gate_uses_cpu, ffn_uses_cpu)) { + fprintf(stderr, "FAIL: bridge_init failed for layer %d\n", layer); + for (int j = 0; j < manifest->n_layer; j++) bridge_release(&bridges[j]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + bridge_mask[layer] = 1; + bridged_layer_count += 1; + } + } else { + if (!bridge_init(&bridges[bridged_layer], [NSString stringWithUTF8String:blob_dir], bridged_layer, bucket, manifest, qkv_input_mode, q_proj_uses_cpu, k_proj_uses_cpu, v_proj_uses_cpu, q_query_uses_cpu, q_gate_uses_cpu, ffn_uses_cpu)) { + fprintf(stderr, "FAIL: bridge_init failed\n"); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + bridge_mask[bridged_layer] = 1; + bridged_layer_count = 1; + } + + int gen_ids[8] = {0}; + float gen_logits[8] = {0.0f}; + int gen_runnerup_ids[8] = {0}; + float gen_runnerup_logits[8] = {0.0f}; + float gen_top2_margins[8] = {0.0f}; + double final_hidden_abs = 0.0; + int seq[128] = {0}; + memcpy(seq, tokens, (size_t)prompt_len * sizeof(int)); + int seq_len = prompt_len; + + for (int step = 0; step < gen_len; step++) { + int next_id = -1; + float next_logit = -INFINITY; + int runnerup_id = -1; + float runnerup_logit = -INFINITY; + if (!compute_next_token_mixed(blob_dir, manifest, bridges, bridge_mask, seq, seq_len, sample_vocab, + lm_head_name_override, &next_id, &next_logit, + &runnerup_id, &runnerup_logit, &final_hidden_abs)) { + fprintf(stderr, "FAIL: mixed top1 compute failed at step %d\n", step); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + gen_ids[step] = next_id; + gen_logits[step] = next_logit; + gen_runnerup_ids[step] = runnerup_id; + gen_runnerup_logits[step] = runnerup_logit; + gen_top2_margins[step] = next_logit - runnerup_logit; + seq[seq_len++] = next_id; + } + + char *decoded_prompt = orion_gpt2_decode(tok, tokens, prompt_len); + char *decoded_gen = orion_gpt2_decode(tok, gen_ids, gen_len); + char *decoded_full = orion_gpt2_decode(tok, seq, seq_len); + + printf("PASS: qwen35 9b decode loop mixed ane cpu smoke\n"); + printf(" prompt=%s\n", prompt); + printf(" prompt_len=%d\n", prompt_len); + printf(" gen_len=%d\n", gen_len); + printf(" bridge_mode=%s\n", mode); + printf(" bucket=%d\n", bucket); + printf(" qkv_input_mode=%s\n", qkv_input_mode_label(qkv_input_mode)); + printf(" q_proj_source=%s\n", q_proj_source_label(q_proj_uses_cpu)); + printf(" q_query_source=%s\n", q_query_source_label(q_query_uses_cpu)); + printf(" k_proj_source=%s\n", k_proj_source_label(k_proj_uses_cpu)); + printf(" q_gate_source=%s\n", q_gate_source_label(q_gate_source, q_gate_uses_cpu, q_gate_cpu_channel_count)); + printf(" ffn_source=%s\n", ffn_source_label(ffn_uses_cpu)); + printf(" q_gate_cpu_channel_count=%d\n", q_gate_cpu_channel_count); + for (int i = 0; i < q_gate_cpu_channel_count; i++) { + printf(" q_gate_cpu_channel_rank=%d channel=%d\n", i + 1, q_gate_cpu_channels[i]); + } + printf(" v_proj_source=%s\n", v_proj_source_label(v_proj_uses_cpu, v_proj_cpu_channel_count)); + printf(" v_proj_cpu_channel_count=%d\n", v_proj_cpu_channel_count); + for (int i = 0; i < v_proj_cpu_channel_count; i++) { + printf(" v_proj_cpu_channel_rank=%d channel=%d\n", i + 1, v_proj_cpu_channels[i]); + } + printf(" bridged_layer=%d\n", bridged_layer); + printf(" bridged_layer_count=%d\n", bridged_layer_count); + printf(" sample_vocab=%d\n", sample_vocab); + printf(" lm_head_name=%s\n", lm_head_name_override ? lm_head_name_override : (manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin")); + for (int i = 0; i < gen_len; i++) { + printf(" gen_token_%d=%d\n", i, gen_ids[i]); + printf(" gen_logit_%d=%.6f\n", i, gen_logits[i]); + printf(" gen_runnerup_token_%d=%d\n", i, gen_runnerup_ids[i]); + printf(" gen_runnerup_logit_%d=%.6f\n", i, gen_runnerup_logits[i]); + printf(" gen_top2_margin_%d=%.6f\n", i, gen_top2_margins[i]); + } + printf(" final_hidden_abs_sum=%.6f\n", final_hidden_abs); + printf(" decoded_prompt=%s\n", decoded_prompt ? decoded_prompt : ""); + printf(" decoded_generated=%s\n", decoded_gen ? decoded_gen : ""); + printf(" decoded_full=%s\n", decoded_full ? decoded_full : ""); + printf(" next_blocker=%s\n", strcmp(mode, "all_full") == 0 ? "evaluate prompt quality and extend beyond full-attention bridge coverage" : "expand bridge from one full-attention layer to all full-attention layers"); + + free(decoded_prompt); + free(decoded_gen); + free(decoded_full); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 0; + } +} diff --git a/tests/test_qwen35_9b_hidden_cache_dump.m b/tests/test_qwen35_9b_hidden_cache_dump.m new file mode 100644 index 0000000..101d6e1 --- /dev/null +++ b/tests/test_qwen35_9b_hidden_cache_dump.m @@ -0,0 +1,147 @@ +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +static NSArray *load_sequences(NSString *path) { + NSString *text = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:nil]; + if (!text) return nil; + NSMutableArray *items = [NSMutableArray array]; + [text enumerateLinesUsingBlock:^(NSString * _Nonnull line, BOOL * _Nonnull stop) { + (void)stop; + if (line.length == 0) return; + NSData *data = [line dataUsingEncoding:NSUTF8StringEncoding]; + NSDictionary *obj = [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; + if ([obj isKindOfClass:[NSDictionary class]]) { + [items addObject:obj]; + } + }]; + return items; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 4) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + NSString *sequencesPath = [NSString stringWithUTF8String:argv[2]]; + NSString *outDir = [NSString stringWithUTF8String:argv[3]]; + + NSArray *sequences = load_sequences(sequencesPath); + if (!sequences || sequences.count == 0) { + fprintf(stderr, "FAIL: sequence file empty or unreadable\n"); + return 1; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + NSMutableOrderedSet *orderedTokens = [NSMutableOrderedSet orderedSet]; + for (NSDictionary *sequence in sequences) { + NSArray *tokenIds = sequence[@"token_ids"]; + if (![tokenIds isKindOfClass:[NSArray class]] || tokenIds.count < 2) continue; + for (NSUInteger idx = 0; idx + 1 < tokenIds.count; idx++) { + [orderedTokens addObject:@([tokenIds[idx] intValue])]; + } + } + + if (orderedTokens.count == 0) { + fprintf(stderr, "FAIL: no input tokens found in sequences\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + + NSError *mkdirError = nil; + [[NSFileManager defaultManager] createDirectoryAtPath:outDir + withIntermediateDirectories:YES + attributes:nil + error:&mkdirError]; + if (mkdirError) { + fprintf(stderr, "FAIL: cannot create cache dir %s\n", outDir.UTF8String); + orion_qwen35_manifest_free(manifest); + return 1; + } + + NSString *binPath = [outDir stringByAppendingPathComponent:@"hidden_cache.bin"]; + NSString *indexPath = [outDir stringByAppendingPathComponent:@"hidden_cache_index.json"]; + FILE *bin = fopen(binPath.UTF8String, "wb"); + if (!bin) { + fprintf(stderr, "FAIL: cannot open hidden cache bin for write\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + + NSMutableArray *tokens = [NSMutableArray arrayWithCapacity:orderedTokens.count]; + NSMutableArray *offsets = [NSMutableArray arrayWithCapacity:orderedTokens.count]; + NSMutableArray *counts = [NSMutableArray arrayWithCapacity:orderedTokens.count]; + long offset = 0; + int dumped = 0; + for (NSNumber *tokenNumber in orderedTokens) { + int tokenId = tokenNumber.intValue; + float *hidden = (float *)calloc((size_t)manifest->d_model, sizeof(float)); + if (!hidden) { + fclose(bin); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, tokenId, hidden)) { + fprintf(stderr, "FAIL: frozen prefix hidden failed for token %d\n", tokenId); + free(hidden); + fclose(bin); + orion_qwen35_manifest_free(manifest); + return 1; + } + + size_t wrote = fwrite(hidden, sizeof(float), (size_t)manifest->d_model, bin); + free(hidden); + if (wrote != (size_t)manifest->d_model) { + fprintf(stderr, "FAIL: hidden cache write short for token %d\n", tokenId); + fclose(bin); + orion_qwen35_manifest_free(manifest); + return 1; + } + + [tokens addObject:tokenNumber]; + [offsets addObject:@(offset)]; + [counts addObject:@(manifest->d_model)]; + offset += (long)manifest->d_model * (long)sizeof(float); + dumped += 1; + } + + fclose(bin); + + NSDictionary *index = @{ + @"status": @"PASS_QWEN35_HIDDEN_CACHE_DUMP", + @"blob_dir": [NSString stringWithUTF8String:blob_dir], + @"sequences_path": sequencesPath, + @"bin_path": binPath, + @"d_model": @(manifest->d_model), + @"tokens": tokens, + @"offsets": offsets, + @"counts": counts, + @"tokens_dumped": @(dumped), + }; + NSData *jsonData = [NSJSONSerialization dataWithJSONObject:index options:NSJSONWritingPrettyPrinted error:nil]; + if (![jsonData writeToFile:indexPath atomically:YES]) { + fprintf(stderr, "FAIL: cannot write hidden cache index\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + + printf("PASS: qwen35 hidden cache dump\n"); + printf(" tokens_dumped=%d\n", dumped); + printf(" d_model=%d\n", manifest->d_model); + printf(" out_dir=%s\n", outDir.UTF8String); + printf(" index_path=%s\n", indexPath.UTF8String); + printf(" bin_path=%s\n", binPath.UTF8String); + + orion_qwen35_manifest_free(manifest); + return 0; + } +} diff --git a/tests/test_qwen35_9b_hybrid_layer_diff.m b/tests/test_qwen35_9b_hybrid_layer_diff.m new file mode 100644 index 0000000..7a1725b --- /dev/null +++ b/tests/test_qwen35_9b_hybrid_layer_diff.m @@ -0,0 +1,3544 @@ +#import +#import +#import +#include +#include +#include +#include +#include +#include +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#import "../model/weight_loader.h" +#import "../tokenizer/gpt2_bpe.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static NSDictionary* load_json(NSString* path) { + NSData* data = [NSData dataWithContentsOfFile:path]; + if (!data) return nil; + return [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; +} + +static int file_exists(const char *path) { + return access(path, F_OK) == 0; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +static double mean_abs_diff(const float *a, const float *b, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)a[i] - (double)b[i]); + return total / (double)n; +} + +static double max_abs_diff(const float *a, const float *b, int n) { + double best = 0.0; + for (int i = 0; i < n; i++) { + double d = fabs((double)a[i] - (double)b[i]); + if (d > best) best = d; + } + return best; +} + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_qproj_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_down_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_down_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static void cpu_linear_batch(const float *x_seq, + int seq_len, + const float *weight, + int in_dim, + int out_dim, + float *out_seq); + +typedef struct { + OrionProgram *prog_q; + OrionProgram *prog_kv; + OrionProgram *prog_ffn; + int layer_idx; + int bucket; + int d_model; + int q_dim; + int kv_dim; + int qkv_input_mode; + int q_uses_cpu_rms; + int kv_uses_cpu_rms; + int q_proj_uses_cpu; + int k_proj_uses_cpu; + int v_proj_uses_cpu; + int q_query_uses_cpu; + int q_gate_uses_cpu; + int ffn_uses_cpu; + int q_gate_cpu_channel_count; + int q_gate_cpu_channels[32]; + int v_proj_cpu_channel_count; + int v_proj_cpu_channels[32]; +} OrionQwen35AneBridge; + +enum { + ORION_QKV_INPUT_MODE_ANE_RMS = 0, + ORION_QKV_INPUT_MODE_CPU_RMS = 1, + ORION_QKV_INPUT_MODE_KV_CPU_RMS = 2, +}; + +static int parse_index_list_local(const char *csv, int *out, int max_count, int limit) { + if (!csv || !*csv || !out || max_count <= 0) return 0; + char *copy = strdup(csv); + if (!copy) return 0; + int count = 0; + char *save = NULL; + for (char *tok = strtok_r(copy, ",", &save); tok && count < max_count; tok = strtok_r(NULL, ",", &save)) { + while (*tok && isspace((unsigned char)*tok)) tok++; + if (!*tok) continue; + char *end = NULL; + long value = strtol(tok, &end, 10); + if (end == tok || value < 0 || value >= limit) continue; + int dup = 0; + for (int i = 0; i < count; i++) { + if (out[i] == (int)value) { + dup = 1; + break; + } + } + if (!dup) out[count++] = (int)value; + } + free(copy); + return count; +} + +static int override_enabled_for_layer(const char *env_name, int layer, int layer_limit) { + const char *csv = getenv(env_name); + if (!csv || !*csv) return 1; + int layers[64] = {0}; + int count = parse_index_list_local(csv, layers, (int)(sizeof(layers) / sizeof(layers[0])), layer_limit); + if (count <= 0) return 1; + for (int i = 0; i < count; i++) { + if (layers[i] == layer) return 1; + } + return 0; +} + +static int load_qkv_input_mode(void) { + const char *mode = getenv("ORION_QKV_INPUT_MODE"); + if (!mode || !*mode) return ORION_QKV_INPUT_MODE_ANE_RMS; + if (strcmp(mode, "cpu_rms") == 0) return ORION_QKV_INPUT_MODE_CPU_RMS; + if (strcmp(mode, "kv_cpu_rms") == 0) return ORION_QKV_INPUT_MODE_KV_CPU_RMS; + return ORION_QKV_INPUT_MODE_ANE_RMS; +} + +static int q_uses_cpu_rms_mode(int qkv_input_mode) { + return qkv_input_mode == ORION_QKV_INPUT_MODE_CPU_RMS; +} + +static int kv_uses_cpu_rms_mode(int qkv_input_mode) { + return qkv_input_mode == ORION_QKV_INPUT_MODE_CPU_RMS || + qkv_input_mode == ORION_QKV_INPUT_MODE_KV_CPU_RMS; +} + +static const char *qkv_input_mode_label(int qkv_input_mode) { + switch (qkv_input_mode) { + case ORION_QKV_INPUT_MODE_CPU_RMS: + return "cpu_rms_linear_only"; + case ORION_QKV_INPUT_MODE_KV_CPU_RMS: + return "kv_cpu_rms_linear_only"; + default: + return "ane_rms_plus_linear"; + } +} + +static int use_cpu_q_proj_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_Q_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_k_proj_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_K_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_query_override(void) { + const char *source = getenv("ORION_Q_QUERY_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_q_gate_override(void) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_cpu_ffn_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_FFN_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int load_q_gate_cpu_channel_preset(const char *source, int *out_channels, int max_channels, int q_dim) { + static const int kSeed2[] = {143, 3994}; + static const int kSeed4[] = {143, 3994, 3768, 1353}; + const int *preset = NULL; + int preset_count = 0; + if (!source || !out_channels || max_channels <= 0) return 0; + if (strcmp(source, "seed2") == 0) { + preset = kSeed2; + preset_count = (int)(sizeof(kSeed2) / sizeof(kSeed2[0])); + } else if (strcmp(source, "seed4") == 0) { + preset = kSeed4; + preset_count = (int)(sizeof(kSeed4) / sizeof(kSeed4[0])); + } + if (!preset) return 0; + int count = 0; + for (int i = 0; i < preset_count && count < max_channels; i++) { + int channel = preset[i]; + if (channel < 0 || channel >= q_dim) continue; + out_channels[count++] = channel; + } + return count; +} + +static int load_q_gate_cpu_channel_override(int *out_channels, int max_channels, int q_dim) { + const char *source = getenv("ORION_Q_GATE_SOURCE"); + int preset_count = load_q_gate_cpu_channel_preset(source, out_channels, max_channels, q_dim); + if (preset_count > 0) return preset_count; + const char *csv = getenv("ORION_Q_GATE_CPU_CHANNELS"); + if (!csv || !*csv) return 0; + return parse_index_list_local(csv, out_channels, max_channels, q_dim); +} + +static int use_cpu_v_proj_override(void) { + const char *full = getenv("ORION_FULL_ATTN_SOURCE"); + if (full && strcmp(full, "cpu") == 0) return 1; + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "cpu") == 0; +} + +static int use_matmul_v_proj_mode(void) { + const char *source = getenv("ORION_V_PROJ_SOURCE"); + return source && strcmp(source, "matmul") == 0; +} + +static int load_v_proj_cpu_channel_override(int *out_channels, int max_channels, int q_dim) { + const char *csv = getenv("ORION_V_PROJ_CPU_CHANNELS"); + if (!csv || !*csv) return 0; + return parse_index_list_local(csv, out_channels, max_channels, q_dim); +} + +static const char *q_proj_source_label(int q_proj_uses_cpu) { + return q_proj_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *q_query_source_label(int q_query_uses_cpu) { + return q_query_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *k_proj_source_label(int k_proj_uses_cpu) { + return k_proj_uses_cpu ? "cpu_linear" : "ane_linear"; +} + +static const char *q_gate_source_label(const char *source, int q_gate_uses_cpu, int partial_channel_count) { + if (q_gate_uses_cpu) return "cpu_linear"; + if (source && strcmp(source, "seed2") == 0) return "ane_linear_seed2"; + if (source && strcmp(source, "seed4") == 0) return "ane_linear_seed4"; + if (partial_channel_count > 0) return "ane_linear_cpu_channels"; + return "ane_linear"; +} + +static const char *ffn_source_label(int ffn_uses_cpu) { + return ffn_uses_cpu ? "cpu" : "ane"; +} + +static const char *v_proj_source_label(int v_proj_uses_cpu, int partial_channel_count) { + if (v_proj_uses_cpu) return "cpu_linear"; + if (partial_channel_count > 0) return "ane_linear_cpu_channels"; + if (use_matmul_v_proj_mode()) return "ane_matmul"; + return "ane_linear"; +} + +static void apply_v_proj_cpu_channel_overrides(float *dst_v_proj_seq, + const float *cpu_v_proj_seq, + int seq_len, + int kv_dim, + int head_dim, + int q_per_kv, + const int *channels, + int channel_count) { + if (!dst_v_proj_seq || !cpu_v_proj_seq || !channels || channel_count <= 0 || seq_len <= 0 || + kv_dim <= 0 || head_dim <= 0 || q_per_kv <= 0) { + return; + } + for (int i = 0; i < channel_count; i++) { + int channel = channels[i]; + int head = channel / head_dim; + int offset = channel % head_dim; + int kv_head = head / q_per_kv; + int kv_channel = kv_head * head_dim + offset; + if (kv_channel < 0 || kv_channel >= kv_dim) continue; + for (int s = 0; s < seq_len; s++) { + dst_v_proj_seq[(size_t)s * kv_dim + kv_channel] = + cpu_v_proj_seq[(size_t)s * kv_dim + kv_channel]; + } + } +} + +static void apply_q_gate_cpu_channel_overrides(float *dst_q_proj_seq, + const float *cpu_q_proj_seq, + int seq_len, + int q_dim, + const int *channels, + int channel_count) { + if (!dst_q_proj_seq || !cpu_q_proj_seq || !channels || channel_count <= 0 || seq_len <= 0 || q_dim <= 0) { + return; + } + for (int i = 0; i < channel_count; i++) { + int channel = channels[i]; + if (channel < 0 || channel >= q_dim) continue; + for (int s = 0; s < seq_len; s++) { + dst_q_proj_seq[(size_t)s * (q_dim * 2) + q_dim + channel] = + cpu_q_proj_seq[(size_t)s * (q_dim * 2) + q_dim + channel]; + } + } +} + +static int load_full_attention_layers(OrionQwen35Manifest *manifest, int *out_layers, int max_layers) { + if (!manifest || !manifest->manifest_path || !out_layers || max_layers <= 0) return 0; + NSString *manifestPath = [NSString stringWithUTF8String:manifest->manifest_path]; + NSDictionary *root = load_json(manifestPath); + NSDictionary *runtime = [root isKindOfClass:[NSDictionary class]] ? root[@"runtime"] : nil; + NSArray *layerTypes = [runtime isKindOfClass:[NSDictionary class]] ? runtime[@"layer_types"] : nil; + if (![layerTypes isKindOfClass:[NSArray class]]) return 0; + + int count = 0; + for (NSInteger i = 0; i < [layerTypes count] && count < max_layers; i++) { + id value = layerTypes[i]; + if ([value isKindOfClass:[NSString class]] && [(NSString *)value isEqualToString:@"full_attention"]) { + out_layers[count++] = (int)i; + } + } + return count; +} + +static void bridge_release(OrionQwen35AneBridge *bridge) { + if (!bridge) return; + if (bridge->prog_q) orion_release_program(bridge->prog_q); + if (bridge->prog_kv) orion_release_program(bridge->prog_kv); + if (bridge->prog_ffn) orion_release_program(bridge->prog_ffn); + memset(bridge, 0, sizeof(*bridge)); +} + +static int bridge_init(OrionQwen35AneBridge *bridge, + NSString *blobDir, + int layer, + int bucket, + OrionQwen35Manifest *manifest, + int qkv_input_mode, + int q_proj_uses_cpu, + int k_proj_uses_cpu, + int v_proj_uses_cpu, + int q_query_uses_cpu, + int q_gate_uses_cpu, + int ffn_uses_cpu) { + memset(bridge, 0, sizeof(*bridge)); + const int q_uses_cpu_rms = q_uses_cpu_rms_mode(qkv_input_mode); + const int kv_uses_cpu_rms = kv_uses_cpu_rms_mode(qkv_input_mode); + OrionModelConfig cfg = { + .n_layer = manifest->n_layer, + .n_head = manifest->n_head, + .n_kv_head = manifest->n_kv_head, + .d_model = manifest->d_model, + .head_dim = manifest->head_dim, + .hidden_dim = manifest->d_ff, + .vocab = manifest->vocab, + .max_seq = manifest->max_seq, + }; + NSString *mil_q = compile_graph( + q_uses_cpu_rms + ? orion_frontend_qwen35_prefill_q_proj_linear_only(layer, bucket, &cfg) + : orion_frontend_qwen35_prefill_q_proj(layer, bucket, &cfg) + ); + NSString *mil_kv = compile_graph( + kv_uses_cpu_rms + ? orion_frontend_qwen35_prefill_kv_proj_linear_only(layer, bucket, &cfg) + : orion_frontend_qwen35_prefill_kv_proj(layer, bucket, &cfg) + ); + NSString *mil_ffn = compile_graph(orion_frontend_qwen35_prefill_ffn(layer, bucket, &cfg)); + if (!mil_q || !mil_kv || !mil_ffn) return 0; + + bridge->prog_q = orion_compile_mil( + mil_q.UTF8String, + q_uses_cpu_rms ? build_qproj_linear_only_wdict(layer, blobDir) : build_qproj_wdict(layer, blobDir), + "qwen35_9b_diff_q" + ); + bridge->prog_kv = orion_compile_mil( + mil_kv.UTF8String, + kv_uses_cpu_rms ? build_kv_linear_only_wdict(layer, blobDir) : build_kv_wdict(layer, blobDir), + "qwen35_9b_diff_kv" + ); + bridge->prog_ffn = orion_compile_mil(mil_ffn.UTF8String, build_ffn_wdict(layer, blobDir), "qwen35_9b_diff_ffn"); + if (!bridge->prog_q || !bridge->prog_kv || !bridge->prog_ffn) { + bridge_release(bridge); + return 0; + } + + bridge->layer_idx = layer; + bridge->bucket = bucket; + bridge->d_model = manifest->d_model; + bridge->q_dim = manifest->n_head * manifest->head_dim; + bridge->kv_dim = manifest->n_kv_head * manifest->head_dim; + bridge->qkv_input_mode = qkv_input_mode; + bridge->q_uses_cpu_rms = q_uses_cpu_rms; + bridge->kv_uses_cpu_rms = kv_uses_cpu_rms; + bridge->q_proj_uses_cpu = q_proj_uses_cpu; + bridge->k_proj_uses_cpu = k_proj_uses_cpu; + bridge->v_proj_uses_cpu = v_proj_uses_cpu; + bridge->q_query_uses_cpu = q_query_uses_cpu; + bridge->q_gate_uses_cpu = q_gate_uses_cpu; + bridge->ffn_uses_cpu = ffn_uses_cpu; + bridge->q_gate_cpu_channel_count = load_q_gate_cpu_channel_override( + bridge->q_gate_cpu_channels, + (int)(sizeof(bridge->q_gate_cpu_channels) / sizeof(bridge->q_gate_cpu_channels[0])), + bridge->q_dim + ); + if (bridge->q_gate_cpu_channel_count > 0 && + !override_enabled_for_layer("ORION_Q_GATE_CPU_CHANNEL_LAYERS", layer, manifest->n_layer)) { + bridge->q_gate_cpu_channel_count = 0; + } + if (bridge->ffn_uses_cpu && + !override_enabled_for_layer("ORION_FFN_CPU_LAYERS", layer, manifest->n_layer)) { + bridge->ffn_uses_cpu = 0; + } + bridge->v_proj_cpu_channel_count = load_v_proj_cpu_channel_override( + bridge->v_proj_cpu_channels, + (int)(sizeof(bridge->v_proj_cpu_channels) / sizeof(bridge->v_proj_cpu_channels[0])), + bridge->q_dim + ); + return 1; +} + +static int load_embeddings(const char *blob_dir, + OrionQwen35Manifest *manifest, + const int *token_ids, + int seq_len, + float *hidden_out) { + const int d_model = manifest->d_model; + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + for (int s = 0; s < seq_len; s++) { + if (!orion_read_blob_row_f32(embed_path, token_ids[s], d_model, hidden_out + s * d_model)) return 0; + } + return 1; +} + +static int apply_cpu_layer(const char *blob_dir, + OrionQwen35Manifest *manifest, + int layer_idx, + const float *hidden_in, + int seq_len, + float *hidden_out) { + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *input_ln = NULL; + float *post_ln = NULL; + float *gate_proj = NULL; + float *up_proj = NULL; + float *down_proj = NULL; + float *normed = NULL; + float *mixer_out = NULL; + float *scratch = NULL; + float *mlp_out = NULL; + + input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) goto fail; + + normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!normed || !mixer_out || !scratch || !mlp_out) goto fail; + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_in + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + if (file_exists(full_q_path)) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + char path_qkv[2048], path_out[2048], path_dt[2048], path_norm[2048], path_conv[2048]; + snprintf(path_qkv, sizeof(path_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + snprintf(path_out, sizeof(path_out), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx); + snprintf(path_dt, sizeof(path_dt), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + snprintf(path_norm, sizeof(path_norm), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx); + snprintf(path_conv, sizeof(path_conv), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + + int qkv_rows = orion_blob_element_count(path_qkv) / d_model; + int value_dim = orion_blob_element_count(path_out) / d_model; + int num_v_heads = orion_blob_element_count(path_dt); + int head_v_dim = orion_blob_element_count(path_norm); + int key_dim = (qkv_rows - value_dim) / 2; + int num_k_heads = num_v_heads; + int head_k_dim = key_dim / num_k_heads; + int conv_kernel = orion_blob_element_count(path_conv) / qkv_rows; + + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) hidden_out[i] = hidden_in[i] + mixer_out[i]; + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_out + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_out[i] += mlp_out[i]; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(mixer_out); free(scratch); free(mlp_out); + return 1; + +fail: + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(mixer_out); free(scratch); free(mlp_out); + return 0; +} + +static int mixed_full_attention_layer(const char *blob_dir, + OrionQwen35Manifest *manifest, + OrionQwen35AneBridge *bridge, + const float *hidden_in, + int seq_len, + float *hidden_out) { + const int layer = bridge->layer_idx; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = bridge->q_dim; + const int kv_dim = bridge->kv_dim; + const int full_cpu_proj_path = + bridge->q_proj_uses_cpu && + bridge->k_proj_uses_cpu && + bridge->v_proj_uses_cpu && + bridge->ffn_uses_cpu; + + float *input_ln = NULL; + float *post_ln = NULL; + float *gate_proj = NULL; + float *up_proj = NULL; + float *down_proj = NULL; + float *o_proj = NULL; + float *q_norm = NULL; + float *k_norm = NULL; + float *q_proj = NULL; + float *k_proj = NULL; + float *v_proj = NULL; + float *normed = NULL; + float *cpu_k_proj_seq = NULL; + float *cpu_v_proj_seq = NULL; + float *q_proj_seq = NULL; + float *k_proj_seq = NULL; + float *v_proj_seq = NULL; + float *attn = NULL; + float *hidden_attn = NULL; + float *post_normed = NULL; + float *mlp = NULL; + IOSurfaceRef ioInQ = NULL; + IOSurfaceRef ioInKV = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + IOSurfaceRef ioFfnIn = NULL; + IOSurfaceRef ioHidden = NULL; + + input_ln = load_exact(blob_dir, layer, "input_layernorm.bin", d_model); + post_ln = load_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + gate_proj = load_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + up_proj = load_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + down_proj = load_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + o_proj = load_exact(blob_dir, layer, "self_attn_o_proj.bin", d_model * q_dim); + q_norm = load_exact(blob_dir, layer, "self_attn_q_norm.bin", head_dim); + k_norm = load_exact(blob_dir, layer, "self_attn_k_norm.bin", head_dim); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0 || full_cpu_proj_path) { + q_proj = load_exact(blob_dir, layer, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + } + if (bridge->k_proj_uses_cpu || full_cpu_proj_path) { + k_proj = load_exact(blob_dir, layer, "self_attn_k_proj.bin", kv_dim * d_model); + } + if (bridge->v_proj_uses_cpu || bridge->v_proj_cpu_channel_count > 0 || full_cpu_proj_path) { + v_proj = load_exact(blob_dir, layer, "self_attn_v_proj.bin", kv_dim * d_model); + } + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj || !o_proj || !q_norm || !k_norm || + ((bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0 || full_cpu_proj_path) && !q_proj) || + ((bridge->k_proj_uses_cpu || full_cpu_proj_path) && !k_proj) || + ((bridge->v_proj_uses_cpu || bridge->v_proj_cpu_channel_count > 0 || full_cpu_proj_path) && !v_proj)) goto fail; + + normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + cpu_k_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + q_proj_seq = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + k_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + v_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + post_normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + mlp = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!normed || !cpu_k_proj_seq || !cpu_v_proj_seq || !q_proj_seq || !k_proj_seq || !v_proj_seq || !attn || !hidden_attn || !post_normed || !mlp) goto fail; + + for (int t = 0; t < seq_len; t++) { + orion_qwen_cpu_rmsnorm(hidden_in + t * d_model, input_ln, d_model, 1e-6f, normed + t * d_model); + } + + if (full_cpu_proj_path) { + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + attn + ); + + for (int i = 0; i < seq_len * d_model; i++) hidden_attn[i] = hidden_in[i] + attn[i]; + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn + s * d_model, post_ln, d_model, 1e-6f, post_normed + s * d_model); + orion_qwen_cpu_swiglu_ffn(post_normed + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_out[i] = hidden_attn[i] + mlp[i]; + + if (input_ln) free(input_ln); + if (post_ln) free(post_ln); + if (gate_proj) free(gate_proj); + if (up_proj) free(up_proj); + if (down_proj) free(down_proj); + if (o_proj) free(o_proj); + if (q_norm) free(q_norm); + if (k_norm) free(k_norm); + if (q_proj) free(q_proj); + if (k_proj) free(k_proj); + if (v_proj) free(v_proj); + if (normed) free(normed); + if (cpu_k_proj_seq) free(cpu_k_proj_seq); + if (cpu_v_proj_seq) free(cpu_v_proj_seq); + if (q_proj_seq) free(q_proj_seq); + if (k_proj_seq) free(k_proj_seq); + if (v_proj_seq) free(v_proj_seq); + if (attn) free(attn); + if (hidden_attn) free(hidden_attn); + if (post_normed) free(post_normed); + if (mlp) free(mlp); + return 1; + } + + const float *q_input = bridge->q_uses_cpu_rms ? normed : hidden_in; + const float *kv_input = bridge->kv_uses_cpu_rms ? normed : hidden_in; + ioInQ = make_cpu_seq_input_surface(q_input, seq_len, bridge->bucket, d_model); + ioInKV = make_cpu_seq_input_surface(kv_input, seq_len, bridge->bucket, d_model); + ioQ = make_f32_surface((q_dim * 2) * bridge->bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bridge->bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bridge->bucket, 0.0f); + + IOSurfaceRef insQ[] = {ioInQ}; + IOSurfaceRef outsQ[] = {ioQ}; + IOSurfaceRef insKV[] = {ioInKV}; + IOSurfaceRef outsKV[] = {ioK, ioV}; + if (!orion_eval(bridge->prog_q, insQ, 1, outsQ, 1) || !orion_eval(bridge->prog_kv, insKV, 1, outsKV, 2)) goto fail_with_surfaces; + + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bridge->bucket, q_proj_seq); + read_ane_surface_prefix(ioK, kv_dim, seq_len, bridge->bucket, k_proj_seq); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bridge->bucket, v_proj_seq); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0) { + float *cpu_q_proj_seq = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + if (!cpu_q_proj_seq) goto fail_with_surfaces; + cpu_linear_batch(normed, seq_len, q_proj, d_model, q_dim * 2, cpu_q_proj_seq); + for (int s = 0; s < seq_len; s++) { + float *dst = q_proj_seq + (size_t)s * (q_dim * 2); + const float *src = cpu_q_proj_seq + (size_t)s * (q_dim * 2); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu) { + memcpy(dst, src, (size_t)q_dim * sizeof(float)); + } + if (bridge->q_proj_uses_cpu || bridge->q_gate_uses_cpu) { + memcpy(dst + q_dim, src + q_dim, (size_t)q_dim * sizeof(float)); + } + } + if (!bridge->q_proj_uses_cpu && !bridge->q_gate_uses_cpu && bridge->q_gate_cpu_channel_count > 0) { + apply_q_gate_cpu_channel_overrides(q_proj_seq, cpu_q_proj_seq, seq_len, q_dim, + bridge->q_gate_cpu_channels, bridge->q_gate_cpu_channel_count); + } + free(cpu_q_proj_seq); + } + if (bridge->k_proj_uses_cpu) { + cpu_linear_batch(normed, seq_len, k_proj, d_model, kv_dim, cpu_k_proj_seq); + memcpy(k_proj_seq, cpu_k_proj_seq, (size_t)seq_len * kv_dim * sizeof(float)); + } + if (bridge->v_proj_uses_cpu || bridge->v_proj_cpu_channel_count > 0) { + cpu_linear_batch(normed, seq_len, v_proj, d_model, kv_dim, cpu_v_proj_seq); + } + if (bridge->v_proj_uses_cpu) { + memcpy(v_proj_seq, cpu_v_proj_seq, (size_t)seq_len * kv_dim * sizeof(float)); + } else if (bridge->v_proj_cpu_channel_count > 0) { + const int q_per_kv = n_head / n_kv_head; + apply_v_proj_cpu_channel_overrides(v_proj_seq, cpu_v_proj_seq, seq_len, kv_dim, head_dim, q_per_kv, + bridge->v_proj_cpu_channels, bridge->v_proj_cpu_channel_count); + } + + orion_qwen_cpu_full_attention_from_projections_with_rope( + q_proj_seq, k_proj_seq, v_proj_seq, seq_len, + o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + attn + ); + + for (int i = 0; i < seq_len * d_model; i++) hidden_attn[i] = hidden_in[i] + attn[i]; + if (bridge->ffn_uses_cpu) { + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn + s * d_model, post_ln, d_model, 1e-6f, post_normed + s * d_model); + orion_qwen_cpu_swiglu_ffn(post_normed + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_out[i] = hidden_attn[i] + mlp[i]; + } else { + ioFfnIn = make_cpu_seq_input_surface(hidden_attn, seq_len, bridge->bucket, d_model); + ioHidden = make_f32_surface(d_model * bridge->bucket, 0.0f); + IOSurfaceRef insFFN[] = {ioFfnIn}; + IOSurfaceRef outsFFN[] = {ioHidden}; + if (!orion_eval(bridge->prog_ffn, insFFN, 1, outsFFN, 1)) goto fail_with_surfaces; + read_ane_surface_prefix(ioHidden, d_model, seq_len, bridge->bucket, hidden_out); + } + + if (ioInQ) CFRelease(ioInQ); + if (ioInKV) CFRelease(ioInKV); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioHidden) CFRelease(ioHidden); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); free(o_proj); free(q_norm); free(k_norm); free(q_proj); free(k_proj); free(v_proj); + free(normed); free(cpu_k_proj_seq); free(cpu_v_proj_seq); free(q_proj_seq); free(k_proj_seq); free(v_proj_seq); free(attn); free(hidden_attn); free(post_normed); free(mlp); + return 1; + +fail_with_surfaces: + if (ioInQ) CFRelease(ioInQ); + if (ioInKV) CFRelease(ioInKV); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioHidden) CFRelease(ioHidden); +fail: + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); free(o_proj); free(q_norm); free(k_norm); free(q_proj); free(k_proj); free(v_proj); + free(normed); free(cpu_k_proj_seq); free(cpu_v_proj_seq); free(q_proj_seq); free(k_proj_seq); free(v_proj_seq); free(attn); free(hidden_attn); free(post_normed); free(mlp); + return 0; +} + +static int apply_hybrid_layer(const char *blob_dir, + OrionQwen35Manifest *manifest, + OrionQwen35AneBridge *bridges, + const unsigned char *bridge_mask, + int layer_idx, + const float *hidden_in, + int seq_len, + float *hidden_out) { + if (bridge_mask && bridge_mask[layer_idx]) { + return mixed_full_attention_layer(blob_dir, manifest, &bridges[layer_idx], hidden_in, seq_len, hidden_out); + } + return apply_cpu_layer(blob_dir, manifest, layer_idx, hidden_in, seq_len, hidden_out); +} + +static void cpu_linear_batch(const float *x_seq, + int seq_len, + const float *weight, + int in_dim, + int out_dim, + float *out_seq) { + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, out_dim, in_dim, + 1.0f, x_seq, in_dim, weight, in_dim, + 0.0f, out_seq, out_dim); +} + +static int apply_cpu_ffn(const float *hidden_attn, + int seq_len, + int d_model, + int d_ff, + const float *post_ln, + const float *gate_proj, + const float *up_proj, + const float *down_proj, + float *hidden_final_out) { + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!scratch || !mlp_out) { + free(scratch); + free(mlp_out); + return 0; + } + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden_final_out[i] = hidden_attn[i] + mlp_out[i]; + free(scratch); + free(mlp_out); + return 1; +} + +static int sampled_topk_logits(const char *blob_dir, + const char *lm_head_name, + const float *hidden, + int d_model, + int sample_vocab, + int *top_id, + float *top_logit) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + int best_id = -1; + float best = -INFINITY; + for (int tok = 0; tok < sample_vocab; tok++) { + if (!orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + if (dot > best) { + best = dot; + best_id = tok; + } + } + + free(row); + *top_id = best_id; + *top_logit = best; + return 1; +} + +static int selected_token_logits(const char *blob_dir, + const char *lm_head_name, + const float *hidden, + int d_model, + const int *token_ids, + int token_count, + float *out_logits) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + for (int tok_idx = 0; tok_idx < token_count; tok_idx++) { + int tok = token_ids[tok_idx]; + if (tok < 0 || !orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + out_logits[tok_idx] = dot; + } + + free(row); + return 1; +} + +#define ORION_GATE_ATTR_REPORT_MAX 64 +#define ORION_GATE_INPUT_ATTR_REPORT_MAX 16 +#define ORION_ATTN_V_ATTR_REPORT_MAX 16 +#define ORION_FFN_DOWN_ATTR_REPORT_MAX 16 +#define ORION_TRACE_DIM_MAX 16 + +typedef struct { + int channel; + double delta_hidden; + double single_pair_gap; +} OrionGateAttrImpact; + +typedef struct { + int dim; + double abs_contrib; + double signed_contrib; + double delta_ffn_rms; +} OrionGateInputAttrImpact; + +typedef struct { + int dim; + double cpu_norm; + double hybrid_norm; + double delta_norm; + double pair_delta; + double abs_contrib; + double signed_contrib; +} OrionAttnVAttrImpact; + +typedef struct { + int dim; + double cpu_norm; + double hybrid_norm; + double delta_norm; + double pair_delta; + double abs_contrib; + double signed_contrib; +} OrionFfnDownAttrImpact; + +static int compare_gate_attr_impact_asc(const void *lhs, const void *rhs) { + const OrionGateAttrImpact *a = (const OrionGateAttrImpact *)lhs; + const OrionGateAttrImpact *b = (const OrionGateAttrImpact *)rhs; + if (a->single_pair_gap < b->single_pair_gap) return -1; + if (a->single_pair_gap > b->single_pair_gap) return 1; + if (a->channel < b->channel) return -1; + if (a->channel > b->channel) return 1; + return 0; +} + +static int compare_gate_input_attr_impact_desc(const void *lhs, const void *rhs) { + const OrionGateInputAttrImpact *a = (const OrionGateInputAttrImpact *)lhs; + const OrionGateInputAttrImpact *b = (const OrionGateInputAttrImpact *)rhs; + if (a->abs_contrib > b->abs_contrib) return -1; + if (a->abs_contrib < b->abs_contrib) return 1; + if (a->dim < b->dim) return -1; + if (a->dim > b->dim) return 1; + return 0; +} + +static int compare_attn_v_attr_impact_desc(const void *lhs, const void *rhs) { + const OrionAttnVAttrImpact *a = (const OrionAttnVAttrImpact *)lhs; + const OrionAttnVAttrImpact *b = (const OrionAttnVAttrImpact *)rhs; + if (a->abs_contrib > b->abs_contrib) return -1; + if (a->abs_contrib < b->abs_contrib) return 1; + if (a->dim < b->dim) return -1; + if (a->dim > b->dim) return 1; + return 0; +} + +static int compare_ffn_down_attr_impact_desc(const void *lhs, const void *rhs) { + const OrionFfnDownAttrImpact *a = (const OrionFfnDownAttrImpact *)lhs; + const OrionFfnDownAttrImpact *b = (const OrionFfnDownAttrImpact *)rhs; + if (a->abs_contrib > b->abs_contrib) return -1; + if (a->abs_contrib < b->abs_contrib) return 1; + if (a->dim < b->dim) return -1; + if (a->dim > b->dim) return 1; + return 0; +} + +static int parse_dim_list(const char *csv, int *out_dims, int max_dims, int dim_limit) { + if (!csv || !out_dims || max_dims <= 0) return 0; + char *copy = strdup(csv); + if (!copy) return 0; + int count = 0; + char *save = NULL; + for (char *tok = strtok_r(copy, ",", &save); tok && count < max_dims; tok = strtok_r(NULL, ",", &save)) { + while (*tok && isspace((unsigned char)*tok)) tok++; + if (!*tok) continue; + char *end = NULL; + long dim = strtol(tok, &end, 10); + if (end == tok) continue; + while (*end && isspace((unsigned char)*end)) end++; + if (*end != '\0') continue; + if (dim < 0 || dim >= dim_limit) continue; + int duplicate = 0; + for (int i = 0; i < count; i++) { + if (out_dims[i] == (int)dim) { + duplicate = 1; + break; + } + } + if (duplicate) continue; + out_dims[count++] = (int)dim; + } + free(copy); + return count; +} + +static int load_pair_delta_row(const char *blob_dir, + const char *lm_head_name, + int d_model, + int candidate_a, + int candidate_b, + float *pair_delta_out) { + if (!pair_delta_out) return 0; + char path[2048]; + snprintf(path, sizeof(path), "%s/model/%s", blob_dir, lm_head_name); + float *row_a = (float *)malloc((size_t)d_model * sizeof(float)); + float *row_b = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row_a || !row_b) { + free(row_a); + free(row_b); + return 0; + } + int ok = orion_read_blob_row_f32(path, candidate_a, d_model, row_a) && + orion_read_blob_row_f32(path, candidate_b, d_model, row_b); + if (ok) { + for (int i = 0; i < d_model; i++) pair_delta_out[i] = row_a[i] - row_b[i]; + } + free(row_a); + free(row_b); + return ok; +} + +static void pair_gap_from_hidden_last_pair_delta(const float *final_norm, + const float *hidden_last, + const float *pair_delta, + int d_model, + float *normed_last_scratch, + double *out_gap) { + orion_qwen_cpu_rmsnorm(hidden_last, final_norm, d_model, 1e-6f, normed_last_scratch); + double dot = 0.0; + for (int i = 0; i < d_model; i++) dot += (double)normed_last_scratch[i] * (double)pair_delta[i]; + *out_gap = dot; +} + +typedef struct { + double cpu_hidden_attn_pair_gap; + double ane_hidden_attn_pair_gap; + double cpu_final_pair_gap; + double ane_final_pair_gap; + int cpu_hidden_attn_pref_token; + int ane_hidden_attn_pref_token; + int cpu_final_pref_token; + int ane_final_pref_token; +} OrionStagePairSensitivity; + +typedef struct { + double final_pair_gap; + int final_pref_token; + int top_id; + float top_logit; +} OrionTailPairReplay; + +typedef struct { + double normed_mean_abs_diff; + double normed_max_abs_diff; + double mixed_linear_mean_abs_diff; + double mixed_linear_max_abs_diff; + double mixed_conv_mean_abs_diff; + double mixed_conv_max_abs_diff; + double query_mean_abs_diff; + double query_max_abs_diff; + double key_mean_abs_diff; + double key_max_abs_diff; + double value_mean_abs_diff; + double value_max_abs_diff; + double z_mean_abs_diff; + double z_max_abs_diff; + double beta_mean_abs_diff; + double beta_max_abs_diff; + double g_mean_abs_diff; + double g_max_abs_diff; + double core_pre_mean_abs_diff; + double core_pre_max_abs_diff; + double core_mean_abs_diff; + double core_max_abs_diff; + double attn_out_mean_abs_diff; + double attn_out_max_abs_diff; + double hidden_attn_mean_abs_diff; + double hidden_attn_max_abs_diff; + double final_mean_abs_diff; + double final_max_abs_diff; + double cpu_hidden_attn_pair_gap; + double hybrid_hidden_attn_pair_gap; + double cpu_final_pair_gap; + double hybrid_final_pair_gap; + int cpu_hidden_attn_pref_token; + int hybrid_hidden_attn_pref_token; + int cpu_final_pref_token; + int hybrid_final_pref_token; +} OrionLinearStageSensitivity; + +typedef struct { + double normed_mean_abs_diff; + double normed_max_abs_diff; + double q_proj_mean_abs_diff; + double q_proj_max_abs_diff; + double q_gate_mean_abs_diff; + double q_gate_max_abs_diff; + double k_proj_mean_abs_diff; + double k_proj_max_abs_diff; + double v_proj_mean_abs_diff; + double v_proj_max_abs_diff; + double attn_gate_sigmoid_mean_abs_diff; + double attn_gate_sigmoid_max_abs_diff; + double attn_context_mean_abs_diff; + double attn_context_max_abs_diff; + double attn_gated_context_mean_abs_diff; + double attn_gated_context_max_abs_diff; + double attn_out_mean_abs_diff; + double attn_out_max_abs_diff; + double attn_qgate_only_mean_abs_diff; + double attn_qgate_only_max_abs_diff; + double attn_v_only_mean_abs_diff; + double attn_v_only_max_abs_diff; + double attn_context_only_mean_abs_diff; + double attn_context_only_max_abs_diff; + double attn_sigmoid_only_mean_abs_diff; + double attn_sigmoid_only_max_abs_diff; + double attn_gated_context_only_mean_abs_diff; + double attn_gated_context_only_max_abs_diff; + double hidden_attn_mean_abs_diff; + double hidden_attn_max_abs_diff; + double ffn_rms_mean_abs_diff; + double ffn_rms_max_abs_diff; + double ffn_gate_mean_abs_diff; + double ffn_gate_max_abs_diff; + double ffn_up_mean_abs_diff; + double ffn_up_max_abs_diff; + double ffn_silu_mean_abs_diff; + double ffn_silu_max_abs_diff; + double ffn_hidden_mean_abs_diff; + double ffn_hidden_max_abs_diff; + double ffn_down_mean_abs_diff; + double ffn_down_max_abs_diff; + double final_mean_abs_diff; + double final_max_abs_diff; + double cpu_hidden_attn_pair_gap; + double hybrid_hidden_attn_pair_gap; + double cpu_res_hybrid_qgate_hidden_attn_pair_gap; + double cpu_res_hybrid_v_hidden_attn_pair_gap; + double cpu_res_hybrid_context_hidden_attn_pair_gap; + double cpu_res_hybrid_sigmoid_hidden_attn_pair_gap; + double cpu_res_hybrid_gated_context_hidden_attn_pair_gap; + double cpu_final_pair_gap; + double hybrid_final_pair_gap; + double cpu_res_hybrid_qgate_final_pair_gap; + double cpu_res_hybrid_v_final_pair_gap; + double cpu_res_hybrid_context_final_pair_gap; + double cpu_res_hybrid_sigmoid_final_pair_gap; + double cpu_res_hybrid_gated_context_final_pair_gap; + int cpu_hidden_attn_pref_token; + int hybrid_hidden_attn_pref_token; + int cpu_res_hybrid_qgate_hidden_attn_pref_token; + int cpu_res_hybrid_v_hidden_attn_pref_token; + int cpu_res_hybrid_context_hidden_attn_pref_token; + int cpu_res_hybrid_sigmoid_hidden_attn_pref_token; + int cpu_res_hybrid_gated_context_hidden_attn_pref_token; + int cpu_final_pref_token; + int hybrid_final_pref_token; + int cpu_res_hybrid_qgate_final_pref_token; + int cpu_res_hybrid_v_final_pref_token; + int cpu_res_hybrid_context_final_pref_token; + int cpu_res_hybrid_sigmoid_final_pref_token; + int cpu_res_hybrid_gated_context_final_pref_token; + double cpu_res_hybrid_down_final_pair_gap; + double hybrid_res_cpu_down_final_pair_gap; + int cpu_res_hybrid_down_final_pref_token; + int hybrid_res_cpu_down_final_pref_token; + double cpu_res_cpu_silu_hybrid_up_final_pair_gap; + double cpu_res_hybrid_silu_cpu_up_final_pair_gap; + int cpu_res_cpu_silu_hybrid_up_final_pref_token; + int cpu_res_hybrid_silu_cpu_up_final_pref_token; + int gate_attr_total_channels; + int gate_attr_report_count; + int gate_attr_first_cumulative_flip_rank; + int gate_attr_flip_channel; + double gate_attr_flip_delta_hidden; + double gate_attr_flip_cumulative_pair_gap; + double gate_attr_all_cumulative_pair_gap; + int gate_input_attr_channel_limit; + int gate_input_attr_count; + int attn_v_attr_count; + int attn_gated_context_attr_count; + int ffn_down_attr_count; + int gate_attr_channels[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_cpu_gate[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_hybrid_gate[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_delta_gate[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_cpu_silu[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_hybrid_silu[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_delta_silu[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_cpu_up[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_delta_hidden[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_single_pair_gap[ORION_GATE_ATTR_REPORT_MAX]; + int gate_attr_single_pref_token[ORION_GATE_ATTR_REPORT_MAX]; + double gate_attr_cumulative_pair_gap[ORION_GATE_ATTR_REPORT_MAX]; + int gate_attr_cumulative_pref_token[ORION_GATE_ATTR_REPORT_MAX]; + int gate_input_attr_dims[ORION_GATE_INPUT_ATTR_REPORT_MAX]; + double gate_input_attr_abs_contrib[ORION_GATE_INPUT_ATTR_REPORT_MAX]; + double gate_input_attr_signed_contrib[ORION_GATE_INPUT_ATTR_REPORT_MAX]; + double gate_input_attr_delta_ffn_rms[ORION_GATE_INPUT_ATTR_REPORT_MAX]; + int attn_v_attr_dims[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_v_attr_cpu_norm[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_v_attr_hybrid_norm[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_v_attr_delta_norm[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_v_attr_pair_delta[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_v_attr_abs_contrib[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_v_attr_signed_contrib[ORION_ATTN_V_ATTR_REPORT_MAX]; + int attn_gated_context_attr_dims[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_gated_context_attr_cpu_norm[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_gated_context_attr_hybrid_norm[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_gated_context_attr_delta_norm[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_gated_context_attr_pair_delta[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_gated_context_attr_abs_contrib[ORION_ATTN_V_ATTR_REPORT_MAX]; + double attn_gated_context_attr_signed_contrib[ORION_ATTN_V_ATTR_REPORT_MAX]; + int ffn_down_attr_dims[ORION_FFN_DOWN_ATTR_REPORT_MAX]; + double ffn_down_attr_cpu_norm[ORION_FFN_DOWN_ATTR_REPORT_MAX]; + double ffn_down_attr_hybrid_norm[ORION_FFN_DOWN_ATTR_REPORT_MAX]; + double ffn_down_attr_delta_norm[ORION_FFN_DOWN_ATTR_REPORT_MAX]; + double ffn_down_attr_pair_delta[ORION_FFN_DOWN_ATTR_REPORT_MAX]; + double ffn_down_attr_abs_contrib[ORION_FFN_DOWN_ATTR_REPORT_MAX]; + double ffn_down_attr_signed_contrib[ORION_FFN_DOWN_ATTR_REPORT_MAX]; +} OrionFullInputStageSensitivity; + +static void l2norm_vec_local(const float *x, int dim, float eps, float *out) { + float sumsq = 0.0f; + for (int i = 0; i < dim; i++) sumsq += x[i] * x[i]; + float inv = 1.0f / sqrtf(sumsq + eps); + for (int i = 0; i < dim; i++) out[i] = x[i] * inv; +} + +static void qwen_rmsnorm_gated_local(const float *x, + const float *gate, + const float *weight, + int dim, + float eps, + float *out) { + float mean_sq = 0.0f; + for (int i = 0; i < dim; i++) mean_sq += x[i] * x[i]; + mean_sq /= (float)dim; + float inv = 1.0f / sqrtf(mean_sq + eps); + for (int i = 0; i < dim; i++) { + float sig = 1.0f / (1.0f + expf(-gate[i])); + out[i] = (x[i] * inv) * weight[i] * (gate[i] * sig); + } +} + +static inline float sigmoid_scalar_local(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +static void silu_vec_local(const float *x, int n, float *out) { + for (int i = 0; i < n; i++) out[i] = x[i] * sigmoid_scalar_local(x[i]); +} + +static void apply_rope_text_inplace_local(float *q, + float *k, + int seq_len, + int n_q_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor) { + int rotary_dim = (int)(head_dim * partial_rotary_factor); + if (rotary_dim > head_dim) rotary_dim = head_dim; + if (rotary_dim % 2 != 0) rotary_dim -= 1; + if (rotary_dim <= 0) return; + + int half_rot = rotary_dim / 2; + float *inv_freq = (float *)malloc((size_t)half_rot * sizeof(float)); + if (!inv_freq) return; + + for (int i = 0; i < half_rot; i++) { + float exponent = (2.0f * (float)i) / (float)rotary_dim; + inv_freq[i] = 1.0f / powf(rope_theta, exponent); + } + + for (int pos = 0; pos < seq_len; pos++) { + for (int h = 0; h < n_q_head; h++) { + float *qh = q + pos * (n_q_head * head_dim) + h * head_dim; + for (int i = 0; i < half_rot; i++) { + float angle = (float)pos * inv_freq[i]; + float c = cosf(angle); + float s = sinf(angle); + float x0 = qh[i]; + float x1 = qh[i + half_rot]; + qh[i] = x0 * c - x1 * s; + qh[i + half_rot] = x1 * c + x0 * s; + } + } + for (int h = 0; h < n_kv_head; h++) { + float *kh = k + pos * (n_kv_head * head_dim) + h * head_dim; + for (int i = 0; i < half_rot; i++) { + float angle = (float)pos * inv_freq[i]; + float c = cosf(angle); + float s = sinf(angle); + float x0 = kh[i]; + float x1 = kh[i + half_rot]; + kh[i] = x0 * c - x1 * s; + kh[i + half_rot] = x1 * c + x0 * s; + } + } + } + + free(inv_freq); +} + +static int capture_attention_gate_context_from_projected_qkv_local(const float *q_proj_out_seq, + const float *k_proj_out_seq, + const float *v_proj_out_seq, + int seq_len, + const float *q_norm, + const float *k_norm, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + float *gate_sigmoid_out, + float *context_out, + float *gated_context_out) { + if (!q_proj_out_seq || !k_proj_out_seq || !v_proj_out_seq || !q_norm || !k_norm || + !gate_sigmoid_out || !context_out || !gated_context_out || + seq_len <= 0 || n_head <= 0 || n_kv_head <= 0 || head_dim <= 0) { + return 0; + } + + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int q_per_kv = n_head / n_kv_head; + const float scale = 1.0f / sqrtf((float)head_dim); + + float *gate_raw = (float *)calloc((size_t)seq_len * q_dim, sizeof(float)); + float *q_normed = (float *)calloc((size_t)seq_len * q_dim, sizeof(float)); + float *k_normed = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + float *q_rope = (float *)calloc((size_t)seq_len * q_dim, sizeof(float)); + float *k_rope = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + float *scores = (float *)calloc((size_t)seq_len * seq_len, sizeof(float)); + float *probs = (float *)calloc((size_t)seq_len * seq_len, sizeof(float)); + float *qh = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + float *kh = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + float *vh = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + float *context_h = (float *)calloc((size_t)seq_len * head_dim, sizeof(float)); + if (!gate_raw || !q_normed || !k_normed || !q_rope || !k_rope || + !scores || !probs || !qh || !kh || !vh || !context_h) { + free(gate_raw); free(q_normed); free(k_normed); free(q_rope); free(k_rope); + free(scores); free(probs); free(qh); free(kh); free(vh); free(context_h); + return 0; + } + + memset(context_out, 0, (size_t)seq_len * q_dim * sizeof(float)); + memset(gated_context_out, 0, (size_t)seq_len * q_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + const float *q_src = q_proj_out_seq + (size_t)s * (q_dim * 2); + memcpy(q_normed + s * q_dim, q_src, (size_t)q_dim * sizeof(float)); + memcpy(gate_raw + s * q_dim, q_src + q_dim, (size_t)q_dim * sizeof(float)); + memcpy(k_normed + s * kv_dim, k_proj_out_seq + (size_t)s * kv_dim, (size_t)kv_dim * sizeof(float)); + for (int i = 0; i < q_dim; i++) { + gate_sigmoid_out[s * q_dim + i] = sigmoid_scalar_local(gate_raw[s * q_dim + i]); + } + } + + for (int s = 0; s < seq_len; s++) { + for (int h = 0; h < n_head; h++) { + float *q_head = q_normed + s * q_dim + h * head_dim; + l2norm_vec_local(q_head, head_dim, 1e-6f, q_head); + for (int i = 0; i < head_dim; i++) q_head[i] *= q_norm[i] * scale; + } + for (int kvh = 0; kvh < n_kv_head; kvh++) { + float *k_head = k_normed + s * kv_dim + kvh * head_dim; + l2norm_vec_local(k_head, head_dim, 1e-6f, k_head); + for (int i = 0; i < head_dim; i++) k_head[i] *= k_norm[i]; + } + } + + memcpy(q_rope, q_normed, (size_t)seq_len * q_dim * sizeof(float)); + memcpy(k_rope, k_normed, (size_t)seq_len * kv_dim * sizeof(float)); + apply_rope_text_inplace_local(q_rope, k_rope, seq_len, n_head, n_kv_head, head_dim, rope_theta, partial_rotary_factor); + + for (int h = 0; h < n_head; h++) { + const int kv_head = h / q_per_kv; + memset(scores, 0, (size_t)seq_len * seq_len * sizeof(float)); + memset(probs, 0, (size_t)seq_len * seq_len * sizeof(float)); + memset(context_h, 0, (size_t)seq_len * head_dim * sizeof(float)); + + for (int s = 0; s < seq_len; s++) { + memcpy(qh + s * head_dim, q_rope + s * q_dim + h * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(kh + s * head_dim, k_rope + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + memcpy(vh + s * head_dim, v_proj_out_seq + s * kv_dim + kv_head * head_dim, (size_t)head_dim * sizeof(float)); + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, seq_len, head_dim, + scale, qh, head_dim, kh, head_dim, + 0.0f, scores, seq_len); + + for (int i = 0; i < seq_len; i++) { + for (int j = i + 1; j < seq_len; j++) scores[i * seq_len + j] = -INFINITY; + float max_val = scores[i * seq_len]; + for (int j = 1; j <= i; j++) { + if (scores[i * seq_len + j] > max_val) max_val = scores[i * seq_len + j]; + } + float sum = 0.0f; + for (int j = 0; j <= i; j++) { + probs[i * seq_len + j] = expf(scores[i * seq_len + j] - max_val); + sum += probs[i * seq_len + j]; + } + for (int j = 0; j <= i; j++) probs[i * seq_len + j] /= sum; + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + seq_len, head_dim, seq_len, + 1.0f, probs, seq_len, vh, head_dim, + 0.0f, context_h, head_dim); + + for (int s = 0; s < seq_len; s++) { + memcpy(context_out + s * q_dim + h * head_dim, + context_h + s * head_dim, + (size_t)head_dim * sizeof(float)); + } + } + + for (int i = 0; i < seq_len * q_dim; i++) { + gated_context_out[i] = context_out[i] * gate_sigmoid_out[i]; + } + + free(gate_raw); free(q_normed); free(k_normed); free(q_rope); free(k_rope); + free(scores); free(probs); free(qh); free(kh); free(vh); free(context_h); + return 1; +} + +static int capture_full_attention_layer_outputs(const float *hidden_in, + int seq_len, + int d_model, + int d_ff, + int n_head, + int n_kv_head, + int head_dim, + float rope_theta, + float partial_rotary_factor, + const float *input_ln, + const float *post_ln, + const float *q_proj, + const float *k_proj, + const float *v_proj, + const float *o_proj, + const float *q_norm, + const float *k_norm, + const float *gate_proj, + const float *up_proj, + const float *down_proj, + float *normed_out, + float *q_proj_out, + float *k_proj_out, + float *v_proj_out, + float *attn_out, + float *hidden_attn_out, + float *ffn_rms_out, + float *ffn_gate_out, + float *ffn_up_out, + float *ffn_silu_out, + float *ffn_hidden_out, + float *ffn_down_out, + float *final_out) { + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int total = seq_len * d_model; + const int total_ff = seq_len * d_ff; + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_in + s * d_model, input_ln, d_model, 1e-6f, normed_out + s * d_model); + } + + cpu_linear_batch(normed_out, seq_len, q_proj, d_model, q_dim * 2, q_proj_out); + cpu_linear_batch(normed_out, seq_len, k_proj, d_model, kv_dim, k_proj_out); + cpu_linear_batch(normed_out, seq_len, v_proj, d_model, kv_dim, v_proj_out); + + orion_qwen_cpu_full_attention_from_projections_with_rope( + q_proj_out, k_proj_out, v_proj_out, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, rope_theta, partial_rotary_factor, + attn_out + ); + + for (int i = 0; i < total; i++) hidden_attn_out[i] = hidden_in[i] + attn_out[i]; + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_attn_out + s * d_model, post_ln, d_model, 1e-6f, ffn_rms_out + s * d_model); + } + + cpu_linear_batch(ffn_rms_out, seq_len, gate_proj, d_model, d_ff, ffn_gate_out); + cpu_linear_batch(ffn_rms_out, seq_len, up_proj, d_model, d_ff, ffn_up_out); + silu_vec_local(ffn_gate_out, total_ff, ffn_silu_out); + for (int i = 0; i < total_ff; i++) ffn_hidden_out[i] = ffn_silu_out[i] * ffn_up_out[i]; + cpu_linear_batch(ffn_hidden_out, seq_len, down_proj, d_ff, d_model, ffn_down_out); + for (int i = 0; i < total; i++) final_out[i] = hidden_attn_out[i] + ffn_down_out[i]; + return 1; +} + +static int capture_linear_layer_outputs(const float *hidden_in, + int seq_len, + int d_model, + int d_ff, + int num_k_heads, + int num_v_heads, + int head_k_dim, + int head_v_dim, + int conv_kernel, + const float *input_ln, + const float *post_ln, + const float *gate_proj, + const float *up_proj, + const float *down_proj, + const float *in_proj_qkv, + const float *in_proj_z, + const float *in_proj_a, + const float *in_proj_b, + const float *conv1d, + const float *dt_bias, + const float *a_log, + const float *norm_weight, + const float *out_proj, + float *normed_out, + float *mixed_linear_out, + float *mixed_conv_out, + float *query_out, + float *key_out, + float *value_out, + float *z_out, + float *beta_out, + float *g_out, + float *core_pre_out, + float *core_out, + float *attn_out, + float *hidden_attn_out, + float *final_out) { + const int key_dim = num_k_heads * head_k_dim; + const int value_dim = num_v_heads * head_v_dim; + const int conv_dim = key_dim * 2 + value_dim; + const float scale = 1.0f / sqrtf((float)head_k_dim); + + float *state = (float *)calloc((size_t)num_v_heads * head_k_dim * head_v_dim, sizeof(float)); + float *q_norm = (float *)calloc((size_t)head_k_dim, sizeof(float)); + float *k_norm = (float *)calloc((size_t)head_k_dim, sizeof(float)); + float *kv_mem = (float *)calloc((size_t)head_v_dim, sizeof(float)); + float *delta = (float *)calloc((size_t)head_v_dim, sizeof(float)); + float *core_head = (float *)calloc((size_t)head_v_dim, sizeof(float)); + if (!state || !q_norm || !k_norm || !kv_mem || !delta || !core_head) { + free(state); + free(q_norm); + free(k_norm); + free(kv_mem); + free(delta); + free(core_head); + return 0; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_in + s * d_model, input_ln, d_model, 1e-6f, normed_out + s * d_model); + } + + cpu_linear_batch(normed_out, seq_len, in_proj_qkv, d_model, conv_dim, mixed_linear_out); + cpu_linear_batch(normed_out, seq_len, in_proj_z, d_model, value_dim, z_out); + cpu_linear_batch(normed_out, seq_len, in_proj_b, d_model, num_v_heads, beta_out); + cpu_linear_batch(normed_out, seq_len, in_proj_a, d_model, num_v_heads, g_out); + + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < conv_dim; c++) { + float sum = 0.0f; + const float *kernel = conv1d + c * conv_kernel; + for (int k = 0; k < conv_kernel; k++) { + int src_t = t - (conv_kernel - 1) + k; + float x = 0.0f; + if (src_t >= 0 && src_t < seq_len) { + x = mixed_linear_out[src_t * conv_dim + c]; + } + sum += kernel[k] * x; + } + float sig = 1.0f / (1.0f + expf(-sum)); + mixed_conv_out[t * conv_dim + c] = sum * sig; + } + + memcpy(query_out + t * key_dim, mixed_conv_out + t * conv_dim, (size_t)key_dim * sizeof(float)); + memcpy(key_out + t * key_dim, mixed_conv_out + t * conv_dim + key_dim, (size_t)key_dim * sizeof(float)); + memcpy(value_out + t * value_dim, mixed_conv_out + t * conv_dim + key_dim * 2, (size_t)value_dim * sizeof(float)); + } + + for (int i = 0; i < seq_len * num_v_heads; i++) { + beta_out[i] = 1.0f / (1.0f + expf(-beta_out[i])); + } + for (int t = 0; t < seq_len; t++) { + for (int h = 0; h < num_v_heads; h++) { + float a = g_out[t * num_v_heads + h]; + float dt = dt_bias[h]; + float al = a_log[h]; + float softplus = (a + dt > 20.0f) ? (a + dt) : log1pf(expf(a + dt)); + g_out[t * num_v_heads + h] = -expf(al) * softplus; + } + } + + for (int t = 0; t < seq_len; t++) { + for (int h = 0; h < num_v_heads; h++) { + const float *q_t = query_out + t * key_dim + h * head_k_dim; + const float *k_t = key_out + t * key_dim + h * head_k_dim; + const float *v_t = value_out + t * value_dim + h * head_v_dim; + const float *z_t = z_out + t * value_dim + h * head_v_dim; + float *state_h = state + h * head_k_dim * head_v_dim; + float *core_pre_t = core_pre_out + t * value_dim + h * head_v_dim; + float *core_t = core_out + t * value_dim + h * head_v_dim; + + l2norm_vec_local(q_t, head_k_dim, 1e-6f, q_norm); + l2norm_vec_local(k_t, head_k_dim, 1e-6f, k_norm); + for (int i = 0; i < head_k_dim; i++) q_norm[i] *= scale; + + float decay = expf(g_out[t * num_v_heads + h]); + float beta_t = beta_out[t * num_v_heads + h]; + for (int i = 0; i < head_k_dim * head_v_dim; i++) state_h[i] *= decay; + + for (int v_idx = 0; v_idx < head_v_dim; v_idx++) { + float sum = 0.0f; + for (int k_idx = 0; k_idx < head_k_dim; k_idx++) { + sum += state_h[k_idx * head_v_dim + v_idx] * k_norm[k_idx]; + } + kv_mem[v_idx] = sum; + delta[v_idx] = (v_t[v_idx] - sum) * beta_t; + } + for (int k_idx = 0; k_idx < head_k_dim; k_idx++) { + for (int v_idx = 0; v_idx < head_v_dim; v_idx++) { + state_h[k_idx * head_v_dim + v_idx] += k_norm[k_idx] * delta[v_idx]; + } + } + for (int v_idx = 0; v_idx < head_v_dim; v_idx++) { + float sum = 0.0f; + for (int k_idx = 0; k_idx < head_k_dim; k_idx++) { + sum += state_h[k_idx * head_v_dim + v_idx] * q_norm[k_idx]; + } + core_pre_t[v_idx] = sum; + } + qwen_rmsnorm_gated_local(core_pre_t, z_t, norm_weight, head_v_dim, 1e-6f, core_t); + } + } + + cpu_linear_batch(core_out, seq_len, out_proj, value_dim, d_model, attn_out); + for (int i = 0; i < seq_len * d_model; i++) hidden_attn_out[i] = hidden_in[i] + attn_out[i]; + if (!apply_cpu_ffn(hidden_attn_out, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, final_out)) { + free(state); + free(q_norm); + free(k_norm); + free(kv_mem); + free(delta); + free(core_head); + return 0; + } + + free(state); + free(q_norm); + free(k_norm); + free(kv_mem); + free(delta); + free(core_head); + return 1; +} + +static int run_linear_input_stage_compare(const char *blob_dir, + OrionQwen35Manifest *manifest, + int layer, + const float *cpu_hidden_in, + const float *hybrid_hidden_in, + int seq_len, + int candidate_a, + int candidate_b, + OrionLinearStageSensitivity *out) { + if (!manifest || !cpu_hidden_in || !hybrid_hidden_in || !out) return 0; + + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + const int pair_ids[2] = {candidate_a, candidate_b}; + + char path_qkv[2048], path_out[2048], path_dt[2048], path_norm[2048], path_conv[2048]; + snprintf(path_qkv, sizeof(path_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer); + snprintf(path_out, sizeof(path_out), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer); + snprintf(path_dt, sizeof(path_dt), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer); + snprintf(path_norm, sizeof(path_norm), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer); + snprintf(path_conv, sizeof(path_conv), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer); + + const int qkv_rows = orion_blob_element_count(path_qkv) / d_model; + const int value_dim = orion_blob_element_count(path_out) / d_model; + const int num_v_heads = orion_blob_element_count(path_dt); + const int head_v_dim = orion_blob_element_count(path_norm); + const int key_dim = (qkv_rows - value_dim) / 2; + const int num_k_heads = num_v_heads; + const int head_k_dim = key_dim / num_k_heads; + const int conv_kernel = orion_blob_element_count(path_conv) / qkv_rows; + const int conv_dim = qkv_rows; + const int total = seq_len * d_model; + const int total_conv = seq_len * conv_dim; + const int total_key = seq_len * key_dim; + const int total_value = seq_len * value_dim; + const int total_heads = seq_len * num_v_heads; + + float *input_ln = load_exact(blob_dir, layer, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + float *in_proj_qkv = load_exact(blob_dir, layer, "linear_attn_in_proj_qkv.bin", qkv_rows * d_model); + float *in_proj_z = load_exact(blob_dir, layer, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer, "linear_attn_conv1d.bin", qkv_rows * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer, "linear_attn_out_proj.bin", d_model * value_dim); + float *final_norm = NULL; + + float *cpu_normed = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_normed = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_mixed_linear = (float *)calloc((size_t)total_conv, sizeof(float)); + float *hybrid_mixed_linear = (float *)calloc((size_t)total_conv, sizeof(float)); + float *cpu_mixed_conv = (float *)calloc((size_t)total_conv, sizeof(float)); + float *hybrid_mixed_conv = (float *)calloc((size_t)total_conv, sizeof(float)); + float *cpu_query = (float *)calloc((size_t)total_key, sizeof(float)); + float *hybrid_query = (float *)calloc((size_t)total_key, sizeof(float)); + float *cpu_key = (float *)calloc((size_t)total_key, sizeof(float)); + float *hybrid_key = (float *)calloc((size_t)total_key, sizeof(float)); + float *cpu_value = (float *)calloc((size_t)total_value, sizeof(float)); + float *hybrid_value = (float *)calloc((size_t)total_value, sizeof(float)); + float *cpu_z = (float *)calloc((size_t)total_value, sizeof(float)); + float *hybrid_z = (float *)calloc((size_t)total_value, sizeof(float)); + float *cpu_beta = (float *)calloc((size_t)total_heads, sizeof(float)); + float *hybrid_beta = (float *)calloc((size_t)total_heads, sizeof(float)); + float *cpu_g = (float *)calloc((size_t)total_heads, sizeof(float)); + float *hybrid_g = (float *)calloc((size_t)total_heads, sizeof(float)); + float *cpu_core_pre = (float *)calloc((size_t)total_value, sizeof(float)); + float *hybrid_core_pre = (float *)calloc((size_t)total_value, sizeof(float)); + float *cpu_core = (float *)calloc((size_t)total_value, sizeof(float)); + float *hybrid_core = (float *)calloc((size_t)total_value, sizeof(float)); + float *cpu_attn_out = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_attn_out = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_final = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *hybrid_last = (float *)calloc((size_t)d_model, sizeof(float)); + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj || !in_proj_qkv || !in_proj_z || + !in_proj_a || !in_proj_b || !conv1d || !dt_bias || !a_log || !norm_weight || !out_proj || + !final_norm || !cpu_normed || !hybrid_normed || !cpu_mixed_linear || !hybrid_mixed_linear || + !cpu_mixed_conv || !hybrid_mixed_conv || !cpu_query || !hybrid_query || !cpu_key || !hybrid_key || + !cpu_value || !hybrid_value || !cpu_z || !hybrid_z || !cpu_beta || !hybrid_beta || !cpu_g || !hybrid_g || + !cpu_core_pre || !hybrid_core_pre || !cpu_core || !hybrid_core || !cpu_attn_out || !hybrid_attn_out || + !cpu_hidden_attn || !hybrid_hidden_attn || !cpu_final || !hybrid_final || !cpu_last || !hybrid_last) { + goto fail; + } + + if (!capture_linear_layer_outputs(cpu_hidden_in, seq_len, d_model, d_ff, num_k_heads, num_v_heads, head_k_dim, head_v_dim, + conv_kernel, input_ln, post_ln, gate_proj, up_proj, down_proj, + in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, conv1d, dt_bias, a_log, norm_weight, out_proj, + cpu_normed, cpu_mixed_linear, cpu_mixed_conv, cpu_query, cpu_key, cpu_value, cpu_z, cpu_beta, cpu_g, + cpu_core_pre, cpu_core, cpu_attn_out, cpu_hidden_attn, cpu_final) || + !capture_linear_layer_outputs(hybrid_hidden_in, seq_len, d_model, d_ff, num_k_heads, num_v_heads, head_k_dim, head_v_dim, + conv_kernel, input_ln, post_ln, gate_proj, up_proj, down_proj, + in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, conv1d, dt_bias, a_log, norm_weight, out_proj, + hybrid_normed, hybrid_mixed_linear, hybrid_mixed_conv, hybrid_query, hybrid_key, hybrid_value, hybrid_z, hybrid_beta, hybrid_g, + hybrid_core_pre, hybrid_core, hybrid_attn_out, hybrid_hidden_attn, hybrid_final)) { + goto fail; + } + + out->normed_mean_abs_diff = mean_abs_diff(cpu_normed, hybrid_normed, total); + out->normed_max_abs_diff = max_abs_diff(cpu_normed, hybrid_normed, total); + out->mixed_linear_mean_abs_diff = mean_abs_diff(cpu_mixed_linear, hybrid_mixed_linear, total_conv); + out->mixed_linear_max_abs_diff = max_abs_diff(cpu_mixed_linear, hybrid_mixed_linear, total_conv); + out->mixed_conv_mean_abs_diff = mean_abs_diff(cpu_mixed_conv, hybrid_mixed_conv, total_conv); + out->mixed_conv_max_abs_diff = max_abs_diff(cpu_mixed_conv, hybrid_mixed_conv, total_conv); + out->query_mean_abs_diff = mean_abs_diff(cpu_query, hybrid_query, total_key); + out->query_max_abs_diff = max_abs_diff(cpu_query, hybrid_query, total_key); + out->key_mean_abs_diff = mean_abs_diff(cpu_key, hybrid_key, total_key); + out->key_max_abs_diff = max_abs_diff(cpu_key, hybrid_key, total_key); + out->value_mean_abs_diff = mean_abs_diff(cpu_value, hybrid_value, total_value); + out->value_max_abs_diff = max_abs_diff(cpu_value, hybrid_value, total_value); + out->z_mean_abs_diff = mean_abs_diff(cpu_z, hybrid_z, total_value); + out->z_max_abs_diff = max_abs_diff(cpu_z, hybrid_z, total_value); + out->beta_mean_abs_diff = mean_abs_diff(cpu_beta, hybrid_beta, total_heads); + out->beta_max_abs_diff = max_abs_diff(cpu_beta, hybrid_beta, total_heads); + out->g_mean_abs_diff = mean_abs_diff(cpu_g, hybrid_g, total_heads); + out->g_max_abs_diff = max_abs_diff(cpu_g, hybrid_g, total_heads); + out->core_pre_mean_abs_diff = mean_abs_diff(cpu_core_pre, hybrid_core_pre, total_value); + out->core_pre_max_abs_diff = max_abs_diff(cpu_core_pre, hybrid_core_pre, total_value); + out->core_mean_abs_diff = mean_abs_diff(cpu_core, hybrid_core, total_value); + out->core_max_abs_diff = max_abs_diff(cpu_core, hybrid_core, total_value); + out->attn_out_mean_abs_diff = mean_abs_diff(cpu_attn_out, hybrid_attn_out, total); + out->attn_out_max_abs_diff = max_abs_diff(cpu_attn_out, hybrid_attn_out, total); + out->hidden_attn_mean_abs_diff = mean_abs_diff(cpu_hidden_attn, hybrid_hidden_attn, total); + out->hidden_attn_max_abs_diff = max_abs_diff(cpu_hidden_attn, hybrid_hidden_attn, total); + out->final_mean_abs_diff = mean_abs_diff(cpu_final, hybrid_final, total); + out->final_max_abs_diff = max_abs_diff(cpu_final, hybrid_final, total); + + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + float cpu_hidden_pair_logits[2] = {0.0f, 0.0f}; + float hybrid_hidden_pair_logits[2] = {0.0f, 0.0f}; + float cpu_final_pair_logits[2] = {0.0f, 0.0f}; + float hybrid_final_pair_logits[2] = {0.0f, 0.0f}; + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_hidden_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, hybrid_hidden_pair_logits)) { + goto fail; + } + orion_qwen_cpu_rmsnorm(cpu_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_final_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, hybrid_final_pair_logits)) { + goto fail; + } + + out->cpu_hidden_attn_pair_gap = (double)cpu_hidden_pair_logits[0] - (double)cpu_hidden_pair_logits[1]; + out->hybrid_hidden_attn_pair_gap = (double)hybrid_hidden_pair_logits[0] - (double)hybrid_hidden_pair_logits[1]; + out->cpu_final_pair_gap = (double)cpu_final_pair_logits[0] - (double)cpu_final_pair_logits[1]; + out->hybrid_final_pair_gap = (double)hybrid_final_pair_logits[0] - (double)hybrid_final_pair_logits[1]; + out->cpu_hidden_attn_pref_token = (out->cpu_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->hybrid_hidden_attn_pref_token = (out->hybrid_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_final_pref_token = (out->cpu_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->hybrid_final_pref_token = (out->hybrid_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(final_norm); + free(cpu_normed); free(hybrid_normed); free(cpu_mixed_linear); free(hybrid_mixed_linear); free(cpu_mixed_conv); free(hybrid_mixed_conv); + free(cpu_query); free(hybrid_query); free(cpu_key); free(hybrid_key); free(cpu_value); free(hybrid_value); free(cpu_z); free(hybrid_z); + free(cpu_beta); free(hybrid_beta); free(cpu_g); free(hybrid_g); free(cpu_core_pre); free(hybrid_core_pre); free(cpu_core); free(hybrid_core); + free(cpu_attn_out); free(hybrid_attn_out); free(cpu_hidden_attn); free(hybrid_hidden_attn); free(cpu_final); free(hybrid_final); + free(cpu_last); free(hybrid_last); + return 1; + +fail: + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(final_norm); + free(cpu_normed); free(hybrid_normed); free(cpu_mixed_linear); free(hybrid_mixed_linear); free(cpu_mixed_conv); free(hybrid_mixed_conv); + free(cpu_query); free(hybrid_query); free(cpu_key); free(hybrid_key); free(cpu_value); free(hybrid_value); free(cpu_z); free(hybrid_z); + free(cpu_beta); free(hybrid_beta); free(cpu_g); free(hybrid_g); free(cpu_core_pre); free(hybrid_core_pre); free(cpu_core); free(hybrid_core); + free(cpu_attn_out); free(hybrid_attn_out); free(cpu_hidden_attn); free(hybrid_hidden_attn); free(cpu_final); free(hybrid_final); + free(cpu_last); free(hybrid_last); + return 0; +} + +static int run_full_input_stage_compare(const char *blob_dir, + OrionQwen35Manifest *manifest, + int layer, + const float *cpu_hidden_in, + const float *hybrid_hidden_in, + int seq_len, + int candidate_a, + int candidate_b, + int gate_attr_report_topk, + OrionFullInputStageSensitivity *out) { + if (!manifest || !cpu_hidden_in || !hybrid_hidden_in || !out) return 0; + + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int total = seq_len * d_model; + const int total_q = seq_len * (q_dim * 2); + const int total_kv = seq_len * kv_dim; + const int total_context = seq_len * q_dim; + const int total_ff = seq_len * d_ff; + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + const int pair_ids[2] = {candidate_a, candidate_b}; + + float *input_ln = load_exact(blob_dir, layer, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + float *q_proj = load_exact(blob_dir, layer, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer, "self_attn_k_norm.bin", head_dim); + float *gate_proj = load_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + float *final_norm = NULL; + + float *cpu_normed = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_normed = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_q = (float *)calloc((size_t)total_q, sizeof(float)); + float *hybrid_q = (float *)calloc((size_t)total_q, sizeof(float)); + float *cpu_k = (float *)calloc((size_t)total_kv, sizeof(float)); + float *hybrid_k = (float *)calloc((size_t)total_kv, sizeof(float)); + float *cpu_v = (float *)calloc((size_t)total_kv, sizeof(float)); + float *hybrid_v = (float *)calloc((size_t)total_kv, sizeof(float)); + float *cpu_gate_sigmoid = (float *)calloc((size_t)total_context, sizeof(float)); + float *hybrid_gate_sigmoid = (float *)calloc((size_t)total_context, sizeof(float)); + float *cpu_context = (float *)calloc((size_t)total_context, sizeof(float)); + float *hybrid_context = (float *)calloc((size_t)total_context, sizeof(float)); + float *cpu_gated_context = (float *)calloc((size_t)total_context, sizeof(float)); + float *hybrid_gated_context = (float *)calloc((size_t)total_context, sizeof(float)); + float *cpu_attn = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_q_hybrid_gate = (float *)calloc((size_t)total_q, sizeof(float)); + float *cpu_hybrid_qgate_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_v_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_context_gated_context = (float *)calloc((size_t)total_context, sizeof(float)); + float *cpu_hybrid_context_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_sigmoid_gated_context = (float *)calloc((size_t)total_context, sizeof(float)); + float *cpu_hybrid_sigmoid_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_gated_context_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_qgate_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_v_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_context_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_sigmoid_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_gated_context_hidden_attn = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_ffn_rms = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_ffn_rms = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_ffn_gate = (float *)calloc((size_t)total_ff, sizeof(float)); + float *hybrid_ffn_gate = (float *)calloc((size_t)total_ff, sizeof(float)); + float *cpu_ffn_up = (float *)calloc((size_t)total_ff, sizeof(float)); + float *hybrid_ffn_up = (float *)calloc((size_t)total_ff, sizeof(float)); + float *cpu_ffn_silu = (float *)calloc((size_t)total_ff, sizeof(float)); + float *hybrid_ffn_silu = (float *)calloc((size_t)total_ff, sizeof(float)); + float *cpu_ffn_hidden = (float *)calloc((size_t)total_ff, sizeof(float)); + float *hybrid_ffn_hidden = (float *)calloc((size_t)total_ff, sizeof(float)); + float *cpu_ffn_down = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_ffn_down = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_final = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_qgate_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_v_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_context_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_sigmoid_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_hybrid_gated_context_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_res_hybrid_down_final = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_res_cpu_down_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_silu_hybrid_up_hidden = (float *)calloc((size_t)total_ff, sizeof(float)); + float *hybrid_silu_cpu_up_hidden = (float *)calloc((size_t)total_ff, sizeof(float)); + float *cpu_res_cpu_silu_hybrid_up_down = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_res_hybrid_silu_cpu_up_down = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_res_cpu_silu_hybrid_up_final = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_res_hybrid_silu_cpu_up_final = (float *)calloc((size_t)total, sizeof(float)); + float *gate_attr_single_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *gate_attr_cumulative_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *gate_attr_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + OrionGateAttrImpact *gate_attr_impacts = (OrionGateAttrImpact *)calloc((size_t)d_ff, sizeof(OrionGateAttrImpact)); + OrionGateInputAttrImpact *gate_input_attr_impacts = (OrionGateInputAttrImpact *)calloc((size_t)d_model, sizeof(OrionGateInputAttrImpact)); + OrionAttnVAttrImpact *attn_v_attr_impacts = (OrionAttnVAttrImpact *)calloc((size_t)d_model, sizeof(OrionAttnVAttrImpact)); + OrionAttnVAttrImpact *attn_gated_context_attr_impacts = (OrionAttnVAttrImpact *)calloc((size_t)d_model, sizeof(OrionAttnVAttrImpact)); + OrionFfnDownAttrImpact *ffn_down_attr_impacts = (OrionFfnDownAttrImpact *)calloc((size_t)d_model, sizeof(OrionFfnDownAttrImpact)); + float *pair_delta = (float *)calloc((size_t)d_model, sizeof(float)); + float *cpu_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *hybrid_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *cpu_v_hidden_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *hybrid_v_hidden_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *cpu_gated_context_hidden_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *hybrid_gated_context_hidden_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *cpu_ffn_down_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *hybrid_ffn_down_norm_last = (float *)calloc((size_t)d_model, sizeof(float)); + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + + if (!input_ln || !post_ln || !q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm || + !gate_proj || !up_proj || !down_proj || !final_norm || + !cpu_normed || !hybrid_normed || !cpu_q || !hybrid_q || !cpu_k || !hybrid_k || + !cpu_v || !hybrid_v || !cpu_gate_sigmoid || !hybrid_gate_sigmoid || !cpu_context || !hybrid_context || + !cpu_gated_context || !hybrid_gated_context || !cpu_attn || !hybrid_attn || + !cpu_q_hybrid_gate || !cpu_hybrid_qgate_attn || !cpu_hybrid_v_attn || + !cpu_hybrid_context_gated_context || !cpu_hybrid_context_attn || !cpu_hybrid_sigmoid_gated_context || + !cpu_hybrid_sigmoid_attn || !cpu_hybrid_gated_context_attn || + !cpu_hidden_attn || !hybrid_hidden_attn || + !cpu_hybrid_qgate_hidden_attn || !cpu_hybrid_v_hidden_attn || + !cpu_hybrid_context_hidden_attn || !cpu_hybrid_sigmoid_hidden_attn || !cpu_hybrid_gated_context_hidden_attn || + !cpu_ffn_rms || !hybrid_ffn_rms || !cpu_ffn_gate || !hybrid_ffn_gate || !cpu_ffn_up || !hybrid_ffn_up || + !cpu_ffn_silu || !hybrid_ffn_silu || !cpu_ffn_hidden || !hybrid_ffn_hidden || + !cpu_ffn_down || !hybrid_ffn_down || !cpu_final || !hybrid_final || + !cpu_hybrid_qgate_final || !cpu_hybrid_v_final || + !cpu_hybrid_context_final || !cpu_hybrid_sigmoid_final || !cpu_hybrid_gated_context_final || + !cpu_res_hybrid_down_final || !hybrid_res_cpu_down_final || + !cpu_silu_hybrid_up_hidden || !hybrid_silu_cpu_up_hidden || + !cpu_res_cpu_silu_hybrid_up_down || !cpu_res_hybrid_silu_cpu_up_down || + !cpu_res_cpu_silu_hybrid_up_final || !cpu_res_hybrid_silu_cpu_up_final || + !gate_attr_single_last || !gate_attr_cumulative_last || !gate_attr_norm_last || !gate_attr_impacts || !gate_input_attr_impacts || !attn_v_attr_impacts || !attn_gated_context_attr_impacts || !ffn_down_attr_impacts || + !pair_delta || + !cpu_last || !hybrid_last || !cpu_v_hidden_norm_last || !hybrid_v_hidden_norm_last || + !cpu_gated_context_hidden_norm_last || !hybrid_gated_context_hidden_norm_last || + !cpu_ffn_down_norm_last || !hybrid_ffn_down_norm_last) { + goto fail; + } + + if (!capture_full_attention_layer_outputs(cpu_hidden_in, seq_len, d_model, d_ff, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + input_ln, post_ln, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + gate_proj, up_proj, down_proj, + cpu_normed, cpu_q, cpu_k, cpu_v, cpu_attn, cpu_hidden_attn, + cpu_ffn_rms, cpu_ffn_gate, cpu_ffn_up, cpu_ffn_silu, cpu_ffn_hidden, + cpu_ffn_down, cpu_final) || + !capture_full_attention_layer_outputs(hybrid_hidden_in, seq_len, d_model, d_ff, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + input_ln, post_ln, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + gate_proj, up_proj, down_proj, + hybrid_normed, hybrid_q, hybrid_k, hybrid_v, hybrid_attn, hybrid_hidden_attn, + hybrid_ffn_rms, hybrid_ffn_gate, hybrid_ffn_up, hybrid_ffn_silu, hybrid_ffn_hidden, + hybrid_ffn_down, hybrid_final)) { + goto fail; + } + + if (!capture_attention_gate_context_from_projected_qkv_local(cpu_q, cpu_k, cpu_v, seq_len, + q_norm, k_norm, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + cpu_gate_sigmoid, cpu_context, cpu_gated_context) || + !capture_attention_gate_context_from_projected_qkv_local(hybrid_q, hybrid_k, hybrid_v, seq_len, + q_norm, k_norm, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + hybrid_gate_sigmoid, hybrid_context, hybrid_gated_context)) { + goto fail; + } + + memcpy(cpu_q_hybrid_gate, cpu_q, (size_t)total_q * sizeof(float)); + for (int s = 0; s < seq_len; s++) { + memcpy(cpu_q_hybrid_gate + (size_t)s * (q_dim * 2) + q_dim, + hybrid_q + (size_t)s * (q_dim * 2) + q_dim, + (size_t)q_dim * sizeof(float)); + } + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q_hybrid_gate, cpu_k, cpu_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + cpu_hybrid_qgate_attn + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, hybrid_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + cpu_hybrid_v_attn + ); + + for (int i = 0; i < total_context; i++) { + cpu_hybrid_context_gated_context[i] = hybrid_context[i] * cpu_gate_sigmoid[i]; + } + cpu_linear_batch(cpu_hybrid_context_gated_context, seq_len, o_proj, q_dim, d_model, cpu_hybrid_context_attn); + for (int i = 0; i < total_context; i++) { + cpu_hybrid_sigmoid_gated_context[i] = cpu_context[i] * hybrid_gate_sigmoid[i]; + } + cpu_linear_batch(cpu_hybrid_sigmoid_gated_context, seq_len, o_proj, q_dim, d_model, cpu_hybrid_sigmoid_attn); + cpu_linear_batch(hybrid_gated_context, seq_len, o_proj, q_dim, d_model, cpu_hybrid_gated_context_attn); + + for (int i = 0; i < total; i++) { + cpu_hybrid_qgate_hidden_attn[i] = cpu_hidden_in[i] + cpu_hybrid_qgate_attn[i]; + cpu_hybrid_v_hidden_attn[i] = cpu_hidden_in[i] + cpu_hybrid_v_attn[i]; + cpu_hybrid_context_hidden_attn[i] = cpu_hidden_in[i] + cpu_hybrid_context_attn[i]; + cpu_hybrid_sigmoid_hidden_attn[i] = cpu_hidden_in[i] + cpu_hybrid_sigmoid_attn[i]; + cpu_hybrid_gated_context_hidden_attn[i] = cpu_hidden_in[i] + cpu_hybrid_gated_context_attn[i]; + } + if (!apply_cpu_ffn(cpu_hybrid_qgate_hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, cpu_hybrid_qgate_final) || + !apply_cpu_ffn(cpu_hybrid_v_hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, cpu_hybrid_v_final) || + !apply_cpu_ffn(cpu_hybrid_context_hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, cpu_hybrid_context_final) || + !apply_cpu_ffn(cpu_hybrid_sigmoid_hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, cpu_hybrid_sigmoid_final) || + !apply_cpu_ffn(cpu_hybrid_gated_context_hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, cpu_hybrid_gated_context_final)) { + goto fail; + } + + out->normed_mean_abs_diff = mean_abs_diff(cpu_normed, hybrid_normed, total); + out->normed_max_abs_diff = max_abs_diff(cpu_normed, hybrid_normed, total); + out->q_proj_mean_abs_diff = mean_abs_diff(cpu_q, hybrid_q, total_q); + out->q_proj_max_abs_diff = max_abs_diff(cpu_q, hybrid_q, total_q); + out->q_gate_mean_abs_diff = mean_abs_diff(cpu_q + q_dim, hybrid_q + q_dim, seq_len * q_dim); + out->q_gate_max_abs_diff = max_abs_diff(cpu_q + q_dim, hybrid_q + q_dim, seq_len * q_dim); + out->k_proj_mean_abs_diff = mean_abs_diff(cpu_k, hybrid_k, total_kv); + out->k_proj_max_abs_diff = max_abs_diff(cpu_k, hybrid_k, total_kv); + out->v_proj_mean_abs_diff = mean_abs_diff(cpu_v, hybrid_v, total_kv); + out->v_proj_max_abs_diff = max_abs_diff(cpu_v, hybrid_v, total_kv); + out->attn_gate_sigmoid_mean_abs_diff = mean_abs_diff(cpu_gate_sigmoid, hybrid_gate_sigmoid, total_context); + out->attn_gate_sigmoid_max_abs_diff = max_abs_diff(cpu_gate_sigmoid, hybrid_gate_sigmoid, total_context); + out->attn_context_mean_abs_diff = mean_abs_diff(cpu_context, hybrid_context, total_context); + out->attn_context_max_abs_diff = max_abs_diff(cpu_context, hybrid_context, total_context); + out->attn_gated_context_mean_abs_diff = mean_abs_diff(cpu_gated_context, hybrid_gated_context, total_context); + out->attn_gated_context_max_abs_diff = max_abs_diff(cpu_gated_context, hybrid_gated_context, total_context); + out->attn_out_mean_abs_diff = mean_abs_diff(cpu_attn, hybrid_attn, total); + out->attn_out_max_abs_diff = max_abs_diff(cpu_attn, hybrid_attn, total); + out->attn_qgate_only_mean_abs_diff = mean_abs_diff(cpu_attn, cpu_hybrid_qgate_attn, total); + out->attn_qgate_only_max_abs_diff = max_abs_diff(cpu_attn, cpu_hybrid_qgate_attn, total); + out->attn_v_only_mean_abs_diff = mean_abs_diff(cpu_attn, cpu_hybrid_v_attn, total); + out->attn_v_only_max_abs_diff = max_abs_diff(cpu_attn, cpu_hybrid_v_attn, total); + out->attn_context_only_mean_abs_diff = mean_abs_diff(cpu_attn, cpu_hybrid_context_attn, total); + out->attn_context_only_max_abs_diff = max_abs_diff(cpu_attn, cpu_hybrid_context_attn, total); + out->attn_sigmoid_only_mean_abs_diff = mean_abs_diff(cpu_attn, cpu_hybrid_sigmoid_attn, total); + out->attn_sigmoid_only_max_abs_diff = max_abs_diff(cpu_attn, cpu_hybrid_sigmoid_attn, total); + out->attn_gated_context_only_mean_abs_diff = mean_abs_diff(cpu_attn, cpu_hybrid_gated_context_attn, total); + out->attn_gated_context_only_max_abs_diff = max_abs_diff(cpu_attn, cpu_hybrid_gated_context_attn, total); + out->hidden_attn_mean_abs_diff = mean_abs_diff(cpu_hidden_attn, hybrid_hidden_attn, total); + out->hidden_attn_max_abs_diff = max_abs_diff(cpu_hidden_attn, hybrid_hidden_attn, total); + out->ffn_rms_mean_abs_diff = mean_abs_diff(cpu_ffn_rms, hybrid_ffn_rms, total); + out->ffn_rms_max_abs_diff = max_abs_diff(cpu_ffn_rms, hybrid_ffn_rms, total); + out->ffn_gate_mean_abs_diff = mean_abs_diff(cpu_ffn_gate, hybrid_ffn_gate, total_ff); + out->ffn_gate_max_abs_diff = max_abs_diff(cpu_ffn_gate, hybrid_ffn_gate, total_ff); + out->ffn_up_mean_abs_diff = mean_abs_diff(cpu_ffn_up, hybrid_ffn_up, total_ff); + out->ffn_up_max_abs_diff = max_abs_diff(cpu_ffn_up, hybrid_ffn_up, total_ff); + out->ffn_silu_mean_abs_diff = mean_abs_diff(cpu_ffn_silu, hybrid_ffn_silu, total_ff); + out->ffn_silu_max_abs_diff = max_abs_diff(cpu_ffn_silu, hybrid_ffn_silu, total_ff); + out->ffn_hidden_mean_abs_diff = mean_abs_diff(cpu_ffn_hidden, hybrid_ffn_hidden, total_ff); + out->ffn_hidden_max_abs_diff = max_abs_diff(cpu_ffn_hidden, hybrid_ffn_hidden, total_ff); + out->ffn_down_mean_abs_diff = mean_abs_diff(cpu_ffn_down, hybrid_ffn_down, total); + out->ffn_down_max_abs_diff = max_abs_diff(cpu_ffn_down, hybrid_ffn_down, total); + out->final_mean_abs_diff = mean_abs_diff(cpu_final, hybrid_final, total); + out->final_max_abs_diff = max_abs_diff(cpu_final, hybrid_final, total); + + for (int i = 0; i < total; i++) { + cpu_res_hybrid_down_final[i] = cpu_hidden_attn[i] + hybrid_ffn_down[i]; + hybrid_res_cpu_down_final[i] = hybrid_hidden_attn[i] + cpu_ffn_down[i]; + } + for (int i = 0; i < total_ff; i++) { + cpu_silu_hybrid_up_hidden[i] = cpu_ffn_silu[i] * hybrid_ffn_up[i]; + hybrid_silu_cpu_up_hidden[i] = hybrid_ffn_silu[i] * cpu_ffn_up[i]; + } + cpu_linear_batch(cpu_silu_hybrid_up_hidden, seq_len, down_proj, d_ff, d_model, cpu_res_cpu_silu_hybrid_up_down); + cpu_linear_batch(hybrid_silu_cpu_up_hidden, seq_len, down_proj, d_ff, d_model, cpu_res_hybrid_silu_cpu_up_down); + for (int i = 0; i < total; i++) { + cpu_res_cpu_silu_hybrid_up_final[i] = cpu_hidden_attn[i] + cpu_res_cpu_silu_hybrid_up_down[i]; + cpu_res_hybrid_silu_cpu_up_final[i] = cpu_hidden_attn[i] + cpu_res_hybrid_silu_cpu_up_down[i]; + } + + float cpu_hidden_logits[2] = {0.0f, 0.0f}; + float hybrid_hidden_logits[2] = {0.0f, 0.0f}; + float cpu_final_logits[2] = {0.0f, 0.0f}; + float hybrid_final_logits[2] = {0.0f, 0.0f}; + float cpu_res_hybrid_down_logits[2] = {0.0f, 0.0f}; + float hybrid_res_cpu_down_logits[2] = {0.0f, 0.0f}; + float cpu_res_cpu_silu_hybrid_up_logits[2] = {0.0f, 0.0f}; + float cpu_res_hybrid_silu_cpu_up_logits[2] = {0.0f, 0.0f}; + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_hidden_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, hybrid_hidden_logits)) { + goto fail; + } + orion_qwen_cpu_rmsnorm(cpu_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_final_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, hybrid_final_logits)) { + goto fail; + } + orion_qwen_cpu_rmsnorm(cpu_res_hybrid_down_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_res_cpu_down_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_res_hybrid_down_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, hybrid_res_cpu_down_logits)) { + goto fail; + } + orion_qwen_cpu_rmsnorm(cpu_res_cpu_silu_hybrid_up_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(cpu_res_hybrid_silu_cpu_up_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_res_cpu_silu_hybrid_up_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, cpu_res_hybrid_silu_cpu_up_logits)) { + goto fail; + } + + out->cpu_hidden_attn_pair_gap = (double)cpu_hidden_logits[0] - (double)cpu_hidden_logits[1]; + out->hybrid_hidden_attn_pair_gap = (double)hybrid_hidden_logits[0] - (double)hybrid_hidden_logits[1]; + out->cpu_final_pair_gap = (double)cpu_final_logits[0] - (double)cpu_final_logits[1]; + out->hybrid_final_pair_gap = (double)hybrid_final_logits[0] - (double)hybrid_final_logits[1]; + out->cpu_res_hybrid_down_final_pair_gap = (double)cpu_res_hybrid_down_logits[0] - (double)cpu_res_hybrid_down_logits[1]; + out->hybrid_res_cpu_down_final_pair_gap = (double)hybrid_res_cpu_down_logits[0] - (double)hybrid_res_cpu_down_logits[1]; + out->cpu_res_cpu_silu_hybrid_up_final_pair_gap = (double)cpu_res_cpu_silu_hybrid_up_logits[0] - (double)cpu_res_cpu_silu_hybrid_up_logits[1]; + out->cpu_res_hybrid_silu_cpu_up_final_pair_gap = (double)cpu_res_hybrid_silu_cpu_up_logits[0] - (double)cpu_res_hybrid_silu_cpu_up_logits[1]; + out->cpu_hidden_attn_pref_token = (out->cpu_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->hybrid_hidden_attn_pref_token = (out->hybrid_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_final_pref_token = (out->cpu_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->hybrid_final_pref_token = (out->hybrid_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_down_final_pref_token = (out->cpu_res_hybrid_down_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->hybrid_res_cpu_down_final_pref_token = (out->hybrid_res_cpu_down_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_cpu_silu_hybrid_up_final_pref_token = (out->cpu_res_cpu_silu_hybrid_up_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_silu_cpu_up_final_pref_token = (out->cpu_res_hybrid_silu_cpu_up_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + + const float *cpu_final_last = cpu_final + (seq_len - 1) * d_model; + const float *cpu_ffn_rms_last = cpu_ffn_rms + (seq_len - 1) * d_model; + const float *hybrid_ffn_rms_last = hybrid_ffn_rms + (seq_len - 1) * d_model; + const float *cpu_ffn_gate_last = cpu_ffn_gate + (seq_len - 1) * d_ff; + const float *hybrid_ffn_gate_last = hybrid_ffn_gate + (seq_len - 1) * d_ff; + const float *cpu_ffn_silu_last = cpu_ffn_silu + (seq_len - 1) * d_ff; + const float *hybrid_ffn_silu_last = hybrid_ffn_silu + (seq_len - 1) * d_ff; + const float *cpu_ffn_up_last = cpu_ffn_up + (seq_len - 1) * d_ff; + if (!load_pair_delta_row(blob_dir, lm_head_name, d_model, candidate_a, candidate_b, pair_delta)) { + goto fail; + } + + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_qgate_hidden_attn + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_qgate_hidden_attn_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_v_hidden_attn + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_v_hidden_attn_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_context_hidden_attn + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_context_hidden_attn_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_sigmoid_hidden_attn + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_sigmoid_hidden_attn_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_gated_context_hidden_attn + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_gated_context_hidden_attn_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_qgate_final + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_qgate_final_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_v_final + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_v_final_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_context_final + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_context_final_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_sigmoid_final + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_sigmoid_final_pair_gap); + pair_gap_from_hidden_last_pair_delta(final_norm, + cpu_hybrid_gated_context_final + (seq_len - 1) * d_model, + pair_delta, d_model, gate_attr_norm_last, + &out->cpu_res_hybrid_gated_context_final_pair_gap); + out->cpu_res_hybrid_qgate_hidden_attn_pref_token = + (out->cpu_res_hybrid_qgate_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_v_hidden_attn_pref_token = + (out->cpu_res_hybrid_v_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_context_hidden_attn_pref_token = + (out->cpu_res_hybrid_context_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_sigmoid_hidden_attn_pref_token = + (out->cpu_res_hybrid_sigmoid_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_gated_context_hidden_attn_pref_token = + (out->cpu_res_hybrid_gated_context_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_qgate_final_pref_token = + (out->cpu_res_hybrid_qgate_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_v_final_pref_token = + (out->cpu_res_hybrid_v_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_context_final_pref_token = + (out->cpu_res_hybrid_context_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_sigmoid_final_pref_token = + (out->cpu_res_hybrid_sigmoid_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_res_hybrid_gated_context_final_pref_token = + (out->cpu_res_hybrid_gated_context_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + + out->attn_v_attr_count = 0; + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_v_hidden_norm_last); + orion_qwen_cpu_rmsnorm(cpu_hybrid_v_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_v_hidden_norm_last); + for (int d = 0; d < d_model; d++) { + double delta_norm = (double)hybrid_v_hidden_norm_last[d] - (double)cpu_v_hidden_norm_last[d]; + double signed_contrib = delta_norm * (double)pair_delta[d]; + attn_v_attr_impacts[d].dim = d; + attn_v_attr_impacts[d].cpu_norm = cpu_v_hidden_norm_last[d]; + attn_v_attr_impacts[d].hybrid_norm = hybrid_v_hidden_norm_last[d]; + attn_v_attr_impacts[d].delta_norm = delta_norm; + attn_v_attr_impacts[d].pair_delta = pair_delta[d]; + attn_v_attr_impacts[d].abs_contrib = fabs(signed_contrib); + attn_v_attr_impacts[d].signed_contrib = signed_contrib; + } + qsort(attn_v_attr_impacts, (size_t)d_model, sizeof(OrionAttnVAttrImpact), compare_attn_v_attr_impact_desc); + const int attn_v_report_count = (d_model < ORION_ATTN_V_ATTR_REPORT_MAX) ? d_model : ORION_ATTN_V_ATTR_REPORT_MAX; + for (int rank = 0; rank < attn_v_report_count; rank++) { + out->attn_v_attr_dims[rank] = attn_v_attr_impacts[rank].dim; + out->attn_v_attr_cpu_norm[rank] = attn_v_attr_impacts[rank].cpu_norm; + out->attn_v_attr_hybrid_norm[rank] = attn_v_attr_impacts[rank].hybrid_norm; + out->attn_v_attr_delta_norm[rank] = attn_v_attr_impacts[rank].delta_norm; + out->attn_v_attr_pair_delta[rank] = attn_v_attr_impacts[rank].pair_delta; + out->attn_v_attr_abs_contrib[rank] = attn_v_attr_impacts[rank].abs_contrib; + out->attn_v_attr_signed_contrib[rank] = attn_v_attr_impacts[rank].signed_contrib; + out->attn_v_attr_count = rank + 1; + } + + out->attn_gated_context_attr_count = 0; + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_gated_context_hidden_norm_last); + orion_qwen_cpu_rmsnorm(cpu_hybrid_gated_context_hidden_attn + (seq_len - 1) * d_model, + final_norm, d_model, 1e-6f, hybrid_gated_context_hidden_norm_last); + for (int d = 0; d < d_model; d++) { + double delta_norm = (double)hybrid_gated_context_hidden_norm_last[d] - (double)cpu_gated_context_hidden_norm_last[d]; + double signed_contrib = delta_norm * (double)pair_delta[d]; + attn_gated_context_attr_impacts[d].dim = d; + attn_gated_context_attr_impacts[d].cpu_norm = cpu_gated_context_hidden_norm_last[d]; + attn_gated_context_attr_impacts[d].hybrid_norm = hybrid_gated_context_hidden_norm_last[d]; + attn_gated_context_attr_impacts[d].delta_norm = delta_norm; + attn_gated_context_attr_impacts[d].pair_delta = pair_delta[d]; + attn_gated_context_attr_impacts[d].abs_contrib = fabs(signed_contrib); + attn_gated_context_attr_impacts[d].signed_contrib = signed_contrib; + } + qsort(attn_gated_context_attr_impacts, (size_t)d_model, sizeof(OrionAttnVAttrImpact), compare_attn_v_attr_impact_desc); + const int attn_gated_context_report_count = + (d_model < ORION_ATTN_V_ATTR_REPORT_MAX) ? d_model : ORION_ATTN_V_ATTR_REPORT_MAX; + for (int rank = 0; rank < attn_gated_context_report_count; rank++) { + out->attn_gated_context_attr_dims[rank] = attn_gated_context_attr_impacts[rank].dim; + out->attn_gated_context_attr_cpu_norm[rank] = attn_gated_context_attr_impacts[rank].cpu_norm; + out->attn_gated_context_attr_hybrid_norm[rank] = attn_gated_context_attr_impacts[rank].hybrid_norm; + out->attn_gated_context_attr_delta_norm[rank] = attn_gated_context_attr_impacts[rank].delta_norm; + out->attn_gated_context_attr_pair_delta[rank] = attn_gated_context_attr_impacts[rank].pair_delta; + out->attn_gated_context_attr_abs_contrib[rank] = attn_gated_context_attr_impacts[rank].abs_contrib; + out->attn_gated_context_attr_signed_contrib[rank] = attn_gated_context_attr_impacts[rank].signed_contrib; + out->attn_gated_context_attr_count = rank + 1; + } + + out->ffn_down_attr_count = 0; + orion_qwen_cpu_rmsnorm(cpu_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_ffn_down_norm_last); + orion_qwen_cpu_rmsnorm(cpu_res_hybrid_down_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_ffn_down_norm_last); + for (int d = 0; d < d_model; d++) { + double delta_norm = (double)hybrid_ffn_down_norm_last[d] - (double)cpu_ffn_down_norm_last[d]; + double signed_contrib = delta_norm * (double)pair_delta[d]; + ffn_down_attr_impacts[d].dim = d; + ffn_down_attr_impacts[d].cpu_norm = cpu_ffn_down_norm_last[d]; + ffn_down_attr_impacts[d].hybrid_norm = hybrid_ffn_down_norm_last[d]; + ffn_down_attr_impacts[d].delta_norm = delta_norm; + ffn_down_attr_impacts[d].pair_delta = pair_delta[d]; + ffn_down_attr_impacts[d].abs_contrib = fabs(signed_contrib); + ffn_down_attr_impacts[d].signed_contrib = signed_contrib; + } + qsort(ffn_down_attr_impacts, (size_t)d_model, sizeof(OrionFfnDownAttrImpact), compare_ffn_down_attr_impact_desc); + const int ffn_down_report_count = (d_model < ORION_FFN_DOWN_ATTR_REPORT_MAX) ? d_model : ORION_FFN_DOWN_ATTR_REPORT_MAX; + for (int rank = 0; rank < ffn_down_report_count; rank++) { + out->ffn_down_attr_dims[rank] = ffn_down_attr_impacts[rank].dim; + out->ffn_down_attr_cpu_norm[rank] = ffn_down_attr_impacts[rank].cpu_norm; + out->ffn_down_attr_hybrid_norm[rank] = ffn_down_attr_impacts[rank].hybrid_norm; + out->ffn_down_attr_delta_norm[rank] = ffn_down_attr_impacts[rank].delta_norm; + out->ffn_down_attr_pair_delta[rank] = ffn_down_attr_impacts[rank].pair_delta; + out->ffn_down_attr_abs_contrib[rank] = ffn_down_attr_impacts[rank].abs_contrib; + out->ffn_down_attr_signed_contrib[rank] = ffn_down_attr_impacts[rank].signed_contrib; + out->ffn_down_attr_count = rank + 1; + } + + if (gate_attr_report_topk < 0) gate_attr_report_topk = 0; + if (gate_attr_report_topk > ORION_GATE_ATTR_REPORT_MAX) gate_attr_report_topk = ORION_GATE_ATTR_REPORT_MAX; + + for (int i = 0; i < d_ff; i++) { + double delta_hidden = ((double)hybrid_ffn_silu_last[i] - (double)cpu_ffn_silu_last[i]) * + (double)cpu_ffn_up_last[i]; + for (int o = 0; o < d_model; o++) { + gate_attr_single_last[o] = cpu_final_last[o] + (float)(delta_hidden * (double)down_proj[o * d_ff + i]); + } + pair_gap_from_hidden_last_pair_delta(final_norm, gate_attr_single_last, pair_delta, d_model, + gate_attr_norm_last, &gate_attr_impacts[i].single_pair_gap); + gate_attr_impacts[i].channel = i; + gate_attr_impacts[i].delta_hidden = delta_hidden; + } + + qsort(gate_attr_impacts, (size_t)d_ff, sizeof(OrionGateAttrImpact), compare_gate_attr_impact_asc); + + memcpy(gate_attr_cumulative_last, cpu_final_last, (size_t)d_model * sizeof(float)); + out->gate_attr_total_channels = d_ff; + out->gate_attr_report_count = 0; + out->gate_attr_first_cumulative_flip_rank = -1; + out->gate_attr_flip_channel = -1; + out->gate_attr_flip_delta_hidden = 0.0; + out->gate_attr_flip_cumulative_pair_gap = INFINITY; + out->gate_attr_all_cumulative_pair_gap = out->cpu_final_pair_gap; + out->gate_input_attr_channel_limit = 0; + out->gate_input_attr_count = 0; + for (int rank = 0; rank < d_ff; rank++) { + const OrionGateAttrImpact impact = gate_attr_impacts[rank]; + for (int o = 0; o < d_model; o++) { + gate_attr_cumulative_last[o] += (float)(impact.delta_hidden * (double)down_proj[o * d_ff + impact.channel]); + } + + double cumulative_pair_gap = 0.0; + pair_gap_from_hidden_last_pair_delta(final_norm, gate_attr_cumulative_last, pair_delta, d_model, + gate_attr_norm_last, &cumulative_pair_gap); + out->gate_attr_all_cumulative_pair_gap = cumulative_pair_gap; + if (rank < gate_attr_report_topk) { + out->gate_attr_channels[rank] = impact.channel; + out->gate_attr_cpu_gate[rank] = cpu_ffn_gate_last[impact.channel]; + out->gate_attr_hybrid_gate[rank] = hybrid_ffn_gate_last[impact.channel]; + out->gate_attr_delta_gate[rank] = (double)hybrid_ffn_gate_last[impact.channel] - (double)cpu_ffn_gate_last[impact.channel]; + out->gate_attr_cpu_silu[rank] = cpu_ffn_silu_last[impact.channel]; + out->gate_attr_hybrid_silu[rank] = hybrid_ffn_silu_last[impact.channel]; + out->gate_attr_delta_silu[rank] = (double)hybrid_ffn_silu_last[impact.channel] - (double)cpu_ffn_silu_last[impact.channel]; + out->gate_attr_cpu_up[rank] = cpu_ffn_up_last[impact.channel]; + out->gate_attr_delta_hidden[rank] = impact.delta_hidden; + out->gate_attr_single_pair_gap[rank] = impact.single_pair_gap; + out->gate_attr_single_pref_token[rank] = (impact.single_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->gate_attr_cumulative_pair_gap[rank] = cumulative_pair_gap; + out->gate_attr_cumulative_pref_token[rank] = (cumulative_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->gate_attr_report_count = rank + 1; + } + if (out->gate_attr_first_cumulative_flip_rank < 0 && cumulative_pair_gap < 0.0) { + out->gate_attr_first_cumulative_flip_rank = rank + 1; + out->gate_attr_flip_channel = impact.channel; + out->gate_attr_flip_delta_hidden = impact.delta_hidden; + out->gate_attr_flip_cumulative_pair_gap = cumulative_pair_gap; + } + } + + int gate_input_channel_limit = out->gate_attr_first_cumulative_flip_rank; + if (gate_input_channel_limit <= 0) gate_input_channel_limit = gate_attr_report_topk; + if (gate_input_channel_limit > d_ff) gate_input_channel_limit = d_ff; + out->gate_input_attr_channel_limit = gate_input_channel_limit; + for (int d = 0; d < d_model; d++) { + double delta_ffn_rms = (double)hybrid_ffn_rms_last[d] - (double)cpu_ffn_rms_last[d]; + double abs_contrib = 0.0; + double signed_contrib = 0.0; + for (int rank = 0; rank < gate_input_channel_limit; rank++) { + const int channel = gate_attr_impacts[rank].channel; + double contrib = delta_ffn_rms * (double)gate_proj[channel * d_model + d]; + abs_contrib += fabs(contrib); + signed_contrib += contrib; + } + gate_input_attr_impacts[d].dim = d; + gate_input_attr_impacts[d].abs_contrib = abs_contrib; + gate_input_attr_impacts[d].signed_contrib = signed_contrib; + gate_input_attr_impacts[d].delta_ffn_rms = delta_ffn_rms; + } + qsort(gate_input_attr_impacts, (size_t)d_model, sizeof(OrionGateInputAttrImpact), compare_gate_input_attr_impact_desc); + const int gate_input_report_count = (d_model < ORION_GATE_INPUT_ATTR_REPORT_MAX) ? d_model : ORION_GATE_INPUT_ATTR_REPORT_MAX; + for (int rank = 0; rank < gate_input_report_count; rank++) { + out->gate_input_attr_dims[rank] = gate_input_attr_impacts[rank].dim; + out->gate_input_attr_abs_contrib[rank] = gate_input_attr_impacts[rank].abs_contrib; + out->gate_input_attr_signed_contrib[rank] = gate_input_attr_impacts[rank].signed_contrib; + out->gate_input_attr_delta_ffn_rms[rank] = gate_input_attr_impacts[rank].delta_ffn_rms; + out->gate_input_attr_count = rank + 1; + } + + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(gate_proj); free(up_proj); free(down_proj); free(final_norm); + free(cpu_normed); free(hybrid_normed); free(cpu_q); free(hybrid_q); free(cpu_k); free(hybrid_k); free(cpu_v); free(hybrid_v); + free(cpu_gate_sigmoid); free(hybrid_gate_sigmoid); free(cpu_context); free(hybrid_context); free(cpu_gated_context); free(hybrid_gated_context); + free(cpu_attn); free(hybrid_attn); free(cpu_q_hybrid_gate); free(cpu_hybrid_qgate_attn); free(cpu_hybrid_v_attn); free(cpu_hybrid_context_gated_context); free(cpu_hybrid_context_attn); free(cpu_hybrid_sigmoid_gated_context); + free(cpu_hybrid_sigmoid_attn); free(cpu_hybrid_gated_context_attn); + free(cpu_hidden_attn); free(hybrid_hidden_attn); free(cpu_hybrid_qgate_hidden_attn); free(cpu_hybrid_v_hidden_attn); free(cpu_hybrid_context_hidden_attn); free(cpu_hybrid_sigmoid_hidden_attn); free(cpu_hybrid_gated_context_hidden_attn); free(cpu_ffn_rms); free(hybrid_ffn_rms); + free(cpu_ffn_gate); free(hybrid_ffn_gate); free(cpu_ffn_up); free(hybrid_ffn_up); free(cpu_ffn_silu); free(hybrid_ffn_silu); + free(cpu_ffn_hidden); free(hybrid_ffn_hidden); free(cpu_ffn_down); free(hybrid_ffn_down); free(cpu_final); free(hybrid_final); + free(cpu_hybrid_qgate_final); free(cpu_hybrid_v_final); free(cpu_hybrid_context_final); free(cpu_hybrid_sigmoid_final); free(cpu_hybrid_gated_context_final); + free(cpu_res_hybrid_down_final); free(hybrid_res_cpu_down_final); + free(cpu_silu_hybrid_up_hidden); free(hybrid_silu_cpu_up_hidden); + free(cpu_res_cpu_silu_hybrid_up_down); free(cpu_res_hybrid_silu_cpu_up_down); + free(cpu_res_cpu_silu_hybrid_up_final); free(cpu_res_hybrid_silu_cpu_up_final); + free(gate_attr_single_last); free(gate_attr_cumulative_last); free(gate_attr_norm_last); free(gate_attr_impacts); free(gate_input_attr_impacts); free(attn_v_attr_impacts); free(attn_gated_context_attr_impacts); free(ffn_down_attr_impacts); free(pair_delta); + free(cpu_last); free(hybrid_last); free(cpu_v_hidden_norm_last); free(hybrid_v_hidden_norm_last); free(cpu_gated_context_hidden_norm_last); free(hybrid_gated_context_hidden_norm_last); free(cpu_ffn_down_norm_last); free(hybrid_ffn_down_norm_last); + return 1; + +fail: + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(gate_proj); free(up_proj); free(down_proj); free(final_norm); + free(cpu_normed); free(hybrid_normed); free(cpu_q); free(hybrid_q); free(cpu_k); free(hybrid_k); free(cpu_v); free(hybrid_v); + free(cpu_gate_sigmoid); free(hybrid_gate_sigmoid); free(cpu_context); free(hybrid_context); free(cpu_gated_context); free(hybrid_gated_context); + free(cpu_attn); free(hybrid_attn); free(cpu_q_hybrid_gate); free(cpu_hybrid_qgate_attn); free(cpu_hybrid_v_attn); free(cpu_hybrid_context_gated_context); free(cpu_hybrid_context_attn); free(cpu_hybrid_sigmoid_gated_context); + free(cpu_hybrid_sigmoid_attn); free(cpu_hybrid_gated_context_attn); + free(cpu_hidden_attn); free(hybrid_hidden_attn); free(cpu_hybrid_qgate_hidden_attn); free(cpu_hybrid_v_hidden_attn); free(cpu_hybrid_context_hidden_attn); free(cpu_hybrid_sigmoid_hidden_attn); free(cpu_hybrid_gated_context_hidden_attn); free(cpu_ffn_rms); free(hybrid_ffn_rms); + free(cpu_ffn_gate); free(hybrid_ffn_gate); free(cpu_ffn_up); free(hybrid_ffn_up); free(cpu_ffn_silu); free(hybrid_ffn_silu); + free(cpu_ffn_hidden); free(hybrid_ffn_hidden); free(cpu_ffn_down); free(hybrid_ffn_down); free(cpu_final); free(hybrid_final); + free(cpu_hybrid_qgate_final); free(cpu_hybrid_v_final); free(cpu_hybrid_context_final); free(cpu_hybrid_sigmoid_final); free(cpu_hybrid_gated_context_final); + free(cpu_res_hybrid_down_final); free(hybrid_res_cpu_down_final); + free(cpu_silu_hybrid_up_hidden); free(hybrid_silu_cpu_up_hidden); + free(cpu_res_cpu_silu_hybrid_up_down); free(cpu_res_hybrid_silu_cpu_up_down); + free(cpu_res_cpu_silu_hybrid_up_final); free(cpu_res_hybrid_silu_cpu_up_final); + free(gate_attr_single_last); free(gate_attr_cumulative_last); free(gate_attr_norm_last); free(gate_attr_impacts); free(gate_input_attr_impacts); free(attn_v_attr_impacts); free(attn_gated_context_attr_impacts); free(ffn_down_attr_impacts); free(pair_delta); + free(cpu_last); free(hybrid_last); free(cpu_v_hidden_norm_last); free(hybrid_v_hidden_norm_last); free(cpu_gated_context_hidden_norm_last); free(hybrid_gated_context_hidden_norm_last); free(cpu_ffn_down_norm_last); free(hybrid_ffn_down_norm_last); + return 0; +} + +static int run_cpu_tail_pair_replay(const char *blob_dir, + OrionQwen35Manifest *manifest, + int start_layer, + const float *hidden_in, + int seq_len, + int sample_vocab, + int candidate_a, + int candidate_b, + OrionTailPairReplay *out) { + if (!manifest || !hidden_in || !out || start_layer < 0 || start_layer >= manifest->n_layer) return 0; + + const int d_model = manifest->d_model; + const int total = seq_len * d_model; + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + const int pair_ids[2] = {candidate_a, candidate_b}; + + float *curr = (float *)calloc((size_t)total, sizeof(float)); + float *next = (float *)calloc((size_t)total, sizeof(float)); + float *final_norm = NULL; + float *last = (float *)calloc((size_t)d_model, sizeof(float)); + if (!curr || !next || !last) { + free(curr); + free(next); + free(last); + return 0; + } + + memcpy(curr, hidden_in, (size_t)total * sizeof(float)); + for (int layer_idx = start_layer; layer_idx < manifest->n_layer; layer_idx++) { + if (!apply_cpu_layer(blob_dir, manifest, layer_idx, curr, seq_len, next)) { + free(curr); + free(next); + free(last); + return 0; + } + float *tmp = curr; + curr = next; + next = tmp; + } + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) { + free(curr); + free(next); + free(last); + return 0; + } + + orion_qwen_cpu_rmsnorm(curr + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, last); + float pair_logits[2] = {0.0f, 0.0f}; + if (!selected_token_logits(blob_dir, lm_head_name, last, d_model, pair_ids, 2, pair_logits) || + !sampled_topk_logits(blob_dir, lm_head_name, last, d_model, sample_vocab, &out->top_id, &out->top_logit)) { + free(curr); + free(next); + free(last); + free(final_norm); + return 0; + } + + out->final_pair_gap = (double)pair_logits[0] - (double)pair_logits[1]; + out->final_pref_token = (out->final_pair_gap >= 0.0) ? candidate_a : candidate_b; + + free(curr); + free(next); + free(last); + free(final_norm); + return 1; +} + +static int run_same_input_stage_pair_compare(const char *blob_dir, + OrionQwen35Manifest *manifest, + OrionQwen35AneBridge *bridge, + const float *hidden_in, + int seq_len, + int candidate_a, + int candidate_b, + OrionStagePairSensitivity *out) { + if (!bridge || !out) return 0; + + const int layer = bridge->layer_idx; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = bridge->q_dim; + const int kv_dim = bridge->kv_dim; + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + const int pair_ids[2] = {candidate_a, candidate_b}; + + float *input_ln = load_exact(blob_dir, layer, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + float *o_proj = load_exact(blob_dir, layer, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer, "self_attn_k_norm.bin", head_dim); + float *final_norm = NULL; + float *normed = NULL; + float *cpu_q = NULL; + float *cpu_k = NULL; + float *cpu_v = NULL; + float *ane_q = NULL; + float *ane_k = NULL; + float *ane_v = NULL; + float *cpu_attn = NULL; + float *ane_attn = NULL; + float *cpu_hidden_attn = NULL; + float *ane_hidden_attn = NULL; + float *cpu_final = NULL; + float *ane_final = NULL; + float *cpu_last = NULL; + float *ane_last = NULL; + float *q_proj_w = NULL; + float *k_proj_w = NULL; + float *v_proj_w = NULL; + IOSurfaceRef ioInQ = NULL; + IOSurfaceRef ioInKV = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + IOSurfaceRef ioFfnIn = NULL; + IOSurfaceRef ioHidden = NULL; + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + + normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + cpu_q = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + cpu_k = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_v = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_q = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + ane_k = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + ane_v = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + cpu_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + ane_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + cpu_hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + ane_hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + cpu_final = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + ane_final = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + cpu_last = (float *)calloc((size_t)d_model, sizeof(float)); + ane_last = (float *)calloc((size_t)d_model, sizeof(float)); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj || !o_proj || !q_norm || !k_norm || + !final_norm || !normed || !cpu_q || !cpu_k || !cpu_v || !ane_q || !ane_k || !ane_v || + !cpu_attn || !ane_attn || !cpu_hidden_attn || !ane_hidden_attn || !cpu_final || !ane_final || + !cpu_last || !ane_last) { + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden_in + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + q_proj_w = load_exact(blob_dir, layer, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + k_proj_w = load_exact(blob_dir, layer, "self_attn_k_proj.bin", kv_dim * d_model); + v_proj_w = load_exact(blob_dir, layer, "self_attn_v_proj.bin", kv_dim * d_model); + if (!q_proj_w || !k_proj_w || !v_proj_w) { + goto fail; + } + cpu_linear_batch(normed, seq_len, q_proj_w, d_model, q_dim * 2, cpu_q); + cpu_linear_batch(normed, seq_len, k_proj_w, d_model, kv_dim, cpu_k); + cpu_linear_batch(normed, seq_len, v_proj_w, d_model, kv_dim, cpu_v); + + { + const float *q_input = bridge->q_uses_cpu_rms ? normed : hidden_in; + const float *kv_input = bridge->kv_uses_cpu_rms ? normed : hidden_in; + ioInQ = make_cpu_seq_input_surface(q_input, seq_len, bridge->bucket, d_model); + ioInKV = make_cpu_seq_input_surface(kv_input, seq_len, bridge->bucket, d_model); + } + ioQ = make_f32_surface((q_dim * 2) * bridge->bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bridge->bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bridge->bucket, 0.0f); + IOSurfaceRef insQ[] = {ioInQ}; + IOSurfaceRef outsQ[] = {ioQ}; + IOSurfaceRef insKV[] = {ioInKV}; + IOSurfaceRef outsKV[] = {ioK, ioV}; + if (!orion_eval(bridge->prog_q, insQ, 1, outsQ, 1) || !orion_eval(bridge->prog_kv, insKV, 1, outsKV, 2)) { + goto fail; + } + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bridge->bucket, ane_q); + read_ane_surface_prefix(ioK, kv_dim, seq_len, bridge->bucket, ane_k); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bridge->bucket, ane_v); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu || bridge->q_gate_uses_cpu || bridge->q_gate_cpu_channel_count > 0) { + float *cpu_q_proj_seq = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + if (!cpu_q_proj_seq) goto fail; + cpu_linear_batch(normed, seq_len, q_proj_w, d_model, q_dim * 2, cpu_q_proj_seq); + for (int s = 0; s < seq_len; s++) { + float *dst = ane_q + (size_t)s * (q_dim * 2); + const float *src = cpu_q_proj_seq + (size_t)s * (q_dim * 2); + if (bridge->q_proj_uses_cpu || bridge->q_query_uses_cpu) { + memcpy(dst, src, (size_t)q_dim * sizeof(float)); + } + if (bridge->q_proj_uses_cpu || bridge->q_gate_uses_cpu) { + memcpy(dst + q_dim, src + q_dim, (size_t)q_dim * sizeof(float)); + } + } + if (!bridge->q_proj_uses_cpu && !bridge->q_gate_uses_cpu && bridge->q_gate_cpu_channel_count > 0) { + apply_q_gate_cpu_channel_overrides(ane_q, cpu_q_proj_seq, seq_len, q_dim, + bridge->q_gate_cpu_channels, bridge->q_gate_cpu_channel_count); + } + free(cpu_q_proj_seq); + } + if (bridge->v_proj_uses_cpu) { + memcpy(ane_v, cpu_v, (size_t)seq_len * kv_dim * sizeof(float)); + } else if (bridge->v_proj_cpu_channel_count > 0) { + apply_v_proj_cpu_channel_overrides(ane_v, cpu_v, seq_len, kv_dim, head_dim, n_head / n_kv_head, + bridge->v_proj_cpu_channels, bridge->v_proj_cpu_channel_count); + } + + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + cpu_attn + ); + orion_qwen_cpu_full_attention_from_projections_with_rope( + ane_q, ane_k, ane_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + ane_attn + ); + + for (int i = 0; i < seq_len * d_model; i++) { + cpu_hidden_attn[i] = hidden_in[i] + cpu_attn[i]; + ane_hidden_attn[i] = hidden_in[i] + ane_attn[i]; + } + + if (!apply_cpu_ffn(cpu_hidden_attn, seq_len, d_model, d_ff, post_ln, gate_proj, up_proj, down_proj, cpu_final)) { + goto fail; + } + ioFfnIn = make_cpu_seq_input_surface(ane_hidden_attn, seq_len, bridge->bucket, d_model); + ioHidden = make_f32_surface(d_model * bridge->bucket, 0.0f); + IOSurfaceRef insFFN[] = {ioFfnIn}; + IOSurfaceRef outsFFN[] = {ioHidden}; + if (!orion_eval(bridge->prog_ffn, insFFN, 1, outsFFN, 1)) { + goto fail; + } + read_ane_surface_prefix(ioHidden, d_model, seq_len, bridge->bucket, ane_final); + + float cpu_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float ane_hidden_attn_pair_logits[2] = {0.0f, 0.0f}; + float cpu_final_pair_logits[2] = {0.0f, 0.0f}; + float ane_final_pair_logits[2] = {0.0f, 0.0f}; + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(ane_hidden_attn + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, ane_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_hidden_attn_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, ane_last, d_model, pair_ids, 2, ane_hidden_attn_pair_logits)) { + goto fail; + } + orion_qwen_cpu_rmsnorm(cpu_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(ane_final + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, ane_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_final_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, ane_last, d_model, pair_ids, 2, ane_final_pair_logits)) { + goto fail; + } + + out->cpu_hidden_attn_pair_gap = (double)cpu_hidden_attn_pair_logits[0] - (double)cpu_hidden_attn_pair_logits[1]; + out->ane_hidden_attn_pair_gap = (double)ane_hidden_attn_pair_logits[0] - (double)ane_hidden_attn_pair_logits[1]; + out->cpu_final_pair_gap = (double)cpu_final_pair_logits[0] - (double)cpu_final_pair_logits[1]; + out->ane_final_pair_gap = (double)ane_final_pair_logits[0] - (double)ane_final_pair_logits[1]; + out->cpu_hidden_attn_pref_token = (out->cpu_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->ane_hidden_attn_pref_token = (out->ane_hidden_attn_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->cpu_final_pref_token = (out->cpu_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + out->ane_final_pref_token = (out->ane_final_pair_gap >= 0.0) ? candidate_a : candidate_b; + + if (ioInQ) CFRelease(ioInQ); + if (ioInKV) CFRelease(ioInKV); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioHidden) CFRelease(ioHidden); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); free(o_proj); free(q_norm); free(k_norm); + free(final_norm); free(normed); free(cpu_q); free(cpu_k); free(cpu_v); free(ane_q); free(ane_k); free(ane_v); + free(cpu_attn); free(ane_attn); free(cpu_hidden_attn); free(ane_hidden_attn); free(cpu_final); free(ane_final); + free(cpu_last); free(ane_last); + free(q_proj_w); free(k_proj_w); free(v_proj_w); + return 1; + +fail: + if (ioInQ) CFRelease(ioInQ); + if (ioInKV) CFRelease(ioInKV); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (ioFfnIn) CFRelease(ioFfnIn); + if (ioHidden) CFRelease(ioHidden); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); free(o_proj); free(q_norm); free(k_norm); + free(final_norm); free(normed); free(cpu_q); free(cpu_k); free(cpu_v); free(ane_q); free(ane_k); free(ane_v); + free(cpu_attn); free(ane_attn); free(cpu_hidden_attn); free(ane_hidden_attn); free(cpu_final); free(ane_final); + free(cpu_last); free(ane_last); + free(q_proj_w); free(k_proj_w); free(v_proj_w); + return 0; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s [single|all_full] [prompt] [sample_vocab] [candidate_a] [candidate_b] [sensitivity_layer]\n", argv[0]); + return 2; + } + if (!orion_ane_init()) { + fprintf(stderr, "FAIL: orion_ane_init failed\n"); + return 3; + } + + const char *blob_dir = argv[1]; + NSString* tokDir = [NSString stringWithUTF8String:argv[2]]; + NSDictionary* meta = load_json([tokDir stringByAppendingPathComponent:@"meta.json"]); + if (!meta) { + fprintf(stderr, "FAIL: missing tokenizer meta.json\n"); + return 1; + } + NSString* regex = meta[@"regex_pattern"]; + NSString* vocabPath = [tokDir stringByAppendingPathComponent:@"vocab.json"]; + NSString* mergesPath = [tokDir stringByAppendingPathComponent:@"merges.txt"]; + OrionGPT2Tokenizer* tok = orion_gpt2_tokenizer_load_with_regex(vocabPath.UTF8String, mergesPath.UTF8String, regex.UTF8String); + if (!tok) { + fprintf(stderr, "FAIL: tokenizer load failed\n"); + return 1; + } + + OrionQwen35Manifest* manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const char *mode = (argc >= 4) ? argv[3] : "single"; + const int qkv_input_mode = load_qkv_input_mode(); + const int q_proj_uses_cpu = use_cpu_q_proj_override(); + const int k_proj_uses_cpu = use_cpu_k_proj_override(); + const int q_query_uses_cpu = use_cpu_q_query_override(); + const int q_gate_uses_cpu = use_cpu_q_gate_override(); + const int ffn_uses_cpu = use_cpu_ffn_override(); + const char *q_gate_source = getenv("ORION_Q_GATE_SOURCE"); + const int v_proj_uses_cpu = use_cpu_v_proj_override(); + int q_gate_cpu_channels[32] = {0}; + const int q_gate_cpu_channel_count = load_q_gate_cpu_channel_override( + q_gate_cpu_channels, + (int)(sizeof(q_gate_cpu_channels) / sizeof(q_gate_cpu_channels[0])), + manifest->n_head * manifest->head_dim + ); + int v_proj_cpu_channels[32] = {0}; + const int v_proj_cpu_channel_count = load_v_proj_cpu_channel_override( + v_proj_cpu_channels, + (int)(sizeof(v_proj_cpu_channels) / sizeof(v_proj_cpu_channels[0])), + manifest->n_head * manifest->head_dim + ); + const char *prompt = (argc >= 5) ? argv[4] : "사진"; + const int sample_vocab = (argc >= 6) ? atoi(argv[5]) : 4096; + const int candidate_a = (argc >= 7) ? atoi(argv[6]) : -1; + const int candidate_b = (argc >= 8) ? atoi(argv[7]) : -1; + const int pair_enabled = (candidate_a >= 0 && candidate_b >= 0); + const int sensitivity_layer = (argc >= 9) ? atoi(argv[8]) : -1; + int gate_attr_report_topk = (argc >= 10) ? atoi(argv[9]) : 32; + const char *trace_dims_csv = (argc >= 11) ? argv[10] : ""; + int trace_dims[ORION_TRACE_DIM_MAX] = {0}; + int trace_dim_count = parse_dim_list(trace_dims_csv, trace_dims, ORION_TRACE_DIM_MAX, manifest->d_model); + if (gate_attr_report_topk < 0) gate_attr_report_topk = 0; + if (gate_attr_report_topk > ORION_GATE_ATTR_REPORT_MAX) gate_attr_report_topk = ORION_GATE_ATTR_REPORT_MAX; + if (strcmp(mode, "single") != 0 && strcmp(mode, "all_full") != 0) { + fprintf(stderr, "FAIL: unsupported mode '%s'\n", mode); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + int token_ids[128] = {0}; + int seq_len = orion_gpt2_encode(tok, prompt, token_ids, 128); + if (seq_len <= 0) { + fprintf(stderr, "FAIL: prompt encode failed\n"); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const int d_model = manifest->d_model; + const int bucket = 32; + const int total = seq_len * d_model; + NSString *blobDir = [NSString stringWithUTF8String:blob_dir]; + + OrionQwen35AneBridge *bridges = calloc((size_t)manifest->n_layer, sizeof(OrionQwen35AneBridge)); + unsigned char *bridge_mask = calloc((size_t)manifest->n_layer, sizeof(unsigned char)); + OrionQwen35AneBridge sensitivity_bridge = {0}; + OrionQwen35AneBridge *sensitivity_bridge_ptr = NULL; + int sensitivity_bridge_owned = 0; + int bridged_layer = 3; + int bridged_layer_count = 0; + if (!bridges || !bridge_mask) { + fprintf(stderr, "FAIL: bridge allocation failed\n"); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + if (strcmp(mode, "all_full") == 0) { + int full_layers[128] = {0}; + int full_count = load_full_attention_layers(manifest, full_layers, 128); + if (full_count <= 0) { + fprintf(stderr, "FAIL: no full_attention layers found\n"); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + bridged_layer = full_layers[0]; + for (int i = 0; i < full_count; i++) { + int layer = full_layers[i]; + if (!bridge_init(&bridges[layer], blobDir, layer, bucket, manifest, qkv_input_mode, q_proj_uses_cpu, k_proj_uses_cpu, v_proj_uses_cpu, q_query_uses_cpu, q_gate_uses_cpu, ffn_uses_cpu)) { + fprintf(stderr, "FAIL: bridge_init failed for layer %d\n", layer); + for (int j = 0; j < manifest->n_layer; j++) bridge_release(&bridges[j]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + bridge_mask[layer] = 1; + bridged_layer_count += 1; + } + } else { + if (!bridge_init(&bridges[bridged_layer], blobDir, bridged_layer, bucket, manifest, qkv_input_mode, q_proj_uses_cpu, k_proj_uses_cpu, v_proj_uses_cpu, q_query_uses_cpu, q_gate_uses_cpu, ffn_uses_cpu)) { + fprintf(stderr, "FAIL: bridge_init failed\n"); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + bridge_mask[bridged_layer] = 1; + bridged_layer_count = 1; + } + + if (pair_enabled && sensitivity_layer >= 0) { + if (sensitivity_layer >= manifest->n_layer) { + fprintf(stderr, "FAIL: sensitivity layer %d out of range\n", sensitivity_layer); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + char sensitivity_q_path[2048]; + snprintf(sensitivity_q_path, sizeof(sensitivity_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, sensitivity_layer); + if (file_exists(sensitivity_q_path)) { + if (bridge_mask[sensitivity_layer]) { + sensitivity_bridge_ptr = &bridges[sensitivity_layer]; + } else { + if (!bridge_init(&sensitivity_bridge, blobDir, sensitivity_layer, bucket, manifest, qkv_input_mode, q_proj_uses_cpu, k_proj_uses_cpu, v_proj_uses_cpu, q_query_uses_cpu, q_gate_uses_cpu, ffn_uses_cpu)) { + fprintf(stderr, "FAIL: sensitivity bridge_init failed for layer %d\n", sensitivity_layer); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + sensitivity_bridge_ptr = &sensitivity_bridge; + sensitivity_bridge_owned = 1; + } + } + } + + float *cpu_curr = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_next = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_curr = (float *)calloc((size_t)total, sizeof(float)); + float *hybrid_next = (float *)calloc((size_t)total, sizeof(float)); + float *sensitivity_cpu_in = (float *)calloc((size_t)total, sizeof(float)); + float *sensitivity_hybrid_in = (float *)calloc((size_t)total, sizeof(float)); + float *cpu_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *hybrid_last = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = NULL; + if (!cpu_curr || !cpu_next || !hybrid_curr || !hybrid_next || !sensitivity_cpu_in || !sensitivity_hybrid_in || + !cpu_last || !hybrid_last) { + fprintf(stderr, "FAIL: hidden allocation failed\n"); + free(cpu_curr); free(cpu_next); free(hybrid_curr); free(hybrid_next); free(sensitivity_cpu_in); free(sensitivity_hybrid_in); free(cpu_last); free(hybrid_last); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + if (!load_embeddings(blob_dir, manifest, token_ids, seq_len, cpu_curr) || + !load_embeddings(blob_dir, manifest, token_ids, seq_len, hybrid_curr)) { + fprintf(stderr, "FAIL: embedding load failed\n"); + goto fail; + } + + printf("TRACE: qwen35 9b hybrid layer diff\n"); + printf(" prompt=%s\n", prompt); + printf(" prompt_len=%d\n", seq_len); + printf(" bridge_mode=%s\n", mode); + printf(" qkv_input_mode=%s\n", qkv_input_mode_label(qkv_input_mode)); + printf(" q_proj_source=%s\n", q_proj_source_label(q_proj_uses_cpu)); + printf(" q_query_source=%s\n", q_query_source_label(q_query_uses_cpu)); + printf(" k_proj_source=%s\n", k_proj_source_label(k_proj_uses_cpu)); + printf(" q_gate_source=%s\n", q_gate_source_label(q_gate_source, q_gate_uses_cpu, q_gate_cpu_channel_count)); + printf(" ffn_source=%s\n", ffn_source_label(ffn_uses_cpu)); + printf(" q_gate_cpu_channel_count=%d\n", q_gate_cpu_channel_count); + for (int i = 0; i < q_gate_cpu_channel_count; i++) { + printf(" q_gate_cpu_channel_rank=%d channel=%d\n", i + 1, q_gate_cpu_channels[i]); + } + printf(" v_proj_source=%s\n", v_proj_source_label(v_proj_uses_cpu, v_proj_cpu_channel_count)); + printf(" v_proj_cpu_channel_count=%d\n", v_proj_cpu_channel_count); + for (int i = 0; i < v_proj_cpu_channel_count; i++) { + printf(" v_proj_cpu_channel_rank=%d channel=%d\n", i + 1, v_proj_cpu_channels[i]); + } + printf(" bridged_layer=%d\n", bridged_layer); + printf(" bridged_layer_count=%d\n", bridged_layer_count); + printf(" sample_vocab=%d\n", sample_vocab); + printf(" candidate_a=%d\n", candidate_a); + printf(" candidate_b=%d\n", candidate_b); + printf(" sensitivity_layer=%d\n", sensitivity_layer); + printf(" gate_attr_report_topk=%d\n", gate_attr_report_topk); + printf(" trace_dim_count=%d\n", trace_dim_count); + for (int i = 0; i < trace_dim_count; i++) { + printf(" trace_dim_config_rank=%d dim=%d\n", i + 1, trace_dims[i]); + } + + int first_diverged_layer = -1; + double first_diverged_max = 0.0; + int worst_layer = -1; + double worst_mean = 0.0; + double worst_max = 0.0; + int first_pair_pref_diverged_layer = -1; + int worst_pair_gap_layer = -1; + double worst_pair_gap_abs_diff = 0.0; + int first_tail_pref_diverged_layer = -1; + int sensitivity_captured = 0; + const char *sensitivity_stage_kind = "none"; + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) { + fprintf(stderr, "FAIL: final_norm load failed\n"); + goto fail; + } + + const char *lm_head_name = manifest->tie_word_embeddings ? "embed_tokens.bin" : "lm_head.bin"; + + for (int layer_idx = 0; layer_idx < manifest->n_layer; layer_idx++) { + if (layer_idx == sensitivity_layer) { + memcpy(sensitivity_cpu_in, cpu_curr, (size_t)total * sizeof(float)); + memcpy(sensitivity_hybrid_in, hybrid_curr, (size_t)total * sizeof(float)); + sensitivity_captured = 1; + } + if (!apply_cpu_layer(blob_dir, manifest, layer_idx, cpu_curr, seq_len, cpu_next)) { + fprintf(stderr, "FAIL: cpu layer apply failed at layer %d\n", layer_idx); + goto fail; + } + if (!apply_hybrid_layer(blob_dir, manifest, bridges, bridge_mask, layer_idx, hybrid_curr, seq_len, hybrid_next)) { + fprintf(stderr, "FAIL: hybrid layer apply failed at layer %d\n", layer_idx); + goto fail; + } + + double mean_diff = mean_abs_diff(cpu_next, hybrid_next, total); + double max_diff = max_abs_diff(cpu_next, hybrid_next, total); + double cpu_abs = abs_sum(cpu_next, total); + double hybrid_abs = abs_sum(hybrid_next, total); + + char full_q_path[2048]; + snprintf(full_q_path, sizeof(full_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + const char *layer_type = file_exists(full_q_path) ? "full_attention" : "linear_attention"; + int bridged = bridge_mask[layer_idx] ? 1 : 0; + + printf("layer=%d type=%s bridged=%d mean_abs_diff=%.6f max_abs_diff=%.6f cpu_abs_sum=%.6f hybrid_abs_sum=%.6f\n", + layer_idx, layer_type, bridged, mean_diff, max_diff, cpu_abs, hybrid_abs); + for (int i = 0; i < trace_dim_count; i++) { + int dim = trace_dims[i]; + const double cpu_val = (double)cpu_next[(seq_len - 1) * d_model + dim]; + const double hybrid_val = (double)hybrid_next[(seq_len - 1) * d_model + dim]; + printf("trace_dim_layer=%d dim=%d cpu=%.6f hybrid=%.6f diff=%.6f\n", + layer_idx, dim, cpu_val, hybrid_val, fabs(cpu_val - hybrid_val)); + } + + if (first_diverged_layer < 0 && max_diff > 1e-6) { + first_diverged_layer = layer_idx; + first_diverged_max = max_diff; + } + if (worst_layer < 0 || max_diff > worst_max) { + worst_layer = layer_idx; + worst_mean = mean_diff; + worst_max = max_diff; + } + + if (pair_enabled) { + float cpu_pair_logits[2] = {0.0f, 0.0f}; + float hybrid_pair_logits[2] = {0.0f, 0.0f}; + const int pair_ids[2] = {candidate_a, candidate_b}; + orion_qwen_cpu_rmsnorm(cpu_next + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_next + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + if (!selected_token_logits(blob_dir, lm_head_name, cpu_last, d_model, pair_ids, 2, cpu_pair_logits) || + !selected_token_logits(blob_dir, lm_head_name, hybrid_last, d_model, pair_ids, 2, hybrid_pair_logits)) { + fprintf(stderr, "FAIL: selected pair logits failed at layer %d\n", layer_idx); + goto fail; + } + + double cpu_pair_gap = (double)cpu_pair_logits[0] - (double)cpu_pair_logits[1]; + double hybrid_pair_gap = (double)hybrid_pair_logits[0] - (double)hybrid_pair_logits[1]; + int cpu_pref_token = (cpu_pair_gap >= 0.0) ? candidate_a : candidate_b; + int hybrid_pref_token = (hybrid_pair_gap >= 0.0) ? candidate_a : candidate_b; + double pair_gap_abs_diff = fabs(cpu_pair_gap - hybrid_pair_gap); + + printf("pair_layer=%d cpu_cand_a_logit=%.6f cpu_cand_b_logit=%.6f cpu_gap=%.6f cpu_pref_token=%d " + "hybrid_cand_a_logit=%.6f hybrid_cand_b_logit=%.6f hybrid_gap=%.6f hybrid_pref_token=%d pair_gap_abs_diff=%.6f\n", + layer_idx, + cpu_pair_logits[0], cpu_pair_logits[1], cpu_pair_gap, cpu_pref_token, + hybrid_pair_logits[0], hybrid_pair_logits[1], hybrid_pair_gap, hybrid_pref_token, + pair_gap_abs_diff); + + if (first_pair_pref_diverged_layer < 0 && cpu_pref_token != hybrid_pref_token) { + first_pair_pref_diverged_layer = layer_idx; + } + if (worst_pair_gap_layer < 0 || pair_gap_abs_diff > worst_pair_gap_abs_diff) { + worst_pair_gap_layer = layer_idx; + worst_pair_gap_abs_diff = pair_gap_abs_diff; + } + } + + if (pair_enabled) { + OrionTailPairReplay cpu_tail = {0}; + OrionTailPairReplay hybrid_tail = {0}; + if (!run_cpu_tail_pair_replay(blob_dir, manifest, layer_idx, cpu_curr, seq_len, sample_vocab, candidate_a, candidate_b, &cpu_tail) || + !run_cpu_tail_pair_replay(blob_dir, manifest, layer_idx, hybrid_curr, seq_len, sample_vocab, candidate_a, candidate_b, &hybrid_tail)) { + fprintf(stderr, "FAIL: cpu tail pair replay failed at layer %d\n", layer_idx); + goto fail; + } + printf("tail_pair_layer=%d cpu_input_final_pair_gap=%.6f hybrid_input_final_pair_gap=%.6f " + "cpu_input_final_pref_token=%d hybrid_input_final_pref_token=%d " + "cpu_input_top_id=%d hybrid_input_top_id=%d\n", + layer_idx, + cpu_tail.final_pair_gap, hybrid_tail.final_pair_gap, + cpu_tail.final_pref_token, hybrid_tail.final_pref_token, + cpu_tail.top_id, hybrid_tail.top_id); + if (first_tail_pref_diverged_layer < 0 && cpu_tail.final_pref_token != hybrid_tail.final_pref_token) { + first_tail_pref_diverged_layer = layer_idx; + } + } + + float *tmp = cpu_curr; cpu_curr = cpu_next; cpu_next = tmp; + tmp = hybrid_curr; hybrid_curr = hybrid_next; hybrid_next = tmp; + } + + double sensitivity_input_mean_abs_diff = 0.0; + double sensitivity_input_max_abs_diff = 0.0; + OrionStagePairSensitivity sensitivity_cpu_input = {0}; + OrionStagePairSensitivity sensitivity_hybrid_input = {0}; + OrionLinearStageSensitivity linear_sensitivity = {0}; + OrionFullInputStageSensitivity full_input_sensitivity = {0}; + if (pair_enabled && sensitivity_layer >= 0) { + if (!sensitivity_captured) { + fprintf(stderr, "FAIL: sensitivity layer %d was not captured\n", sensitivity_layer); + goto fail; + } + sensitivity_input_mean_abs_diff = mean_abs_diff(sensitivity_cpu_in, sensitivity_hybrid_in, total); + sensitivity_input_max_abs_diff = max_abs_diff(sensitivity_cpu_in, sensitivity_hybrid_in, total); + char sensitivity_q_path[2048]; + snprintf(sensitivity_q_path, sizeof(sensitivity_q_path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, sensitivity_layer); + if (file_exists(sensitivity_q_path)) { + sensitivity_stage_kind = "full_attention"; + if (!sensitivity_bridge_ptr || + !run_same_input_stage_pair_compare(blob_dir, manifest, sensitivity_bridge_ptr, sensitivity_cpu_in, seq_len, candidate_a, candidate_b, &sensitivity_cpu_input) || + !run_same_input_stage_pair_compare(blob_dir, manifest, sensitivity_bridge_ptr, sensitivity_hybrid_in, seq_len, candidate_a, candidate_b, &sensitivity_hybrid_input) || + !run_full_input_stage_compare(blob_dir, manifest, sensitivity_layer, sensitivity_cpu_in, sensitivity_hybrid_in, + seq_len, candidate_a, candidate_b, gate_attr_report_topk, &full_input_sensitivity)) { + fprintf(stderr, "FAIL: sensitivity stage pair compare failed at layer %d\n", sensitivity_layer); + goto fail; + } + } else { + sensitivity_stage_kind = "linear_attention"; + if (!run_linear_input_stage_compare(blob_dir, manifest, sensitivity_layer, sensitivity_cpu_in, sensitivity_hybrid_in, + seq_len, candidate_a, candidate_b, &linear_sensitivity)) { + fprintf(stderr, "FAIL: linear sensitivity stage compare failed at layer %d\n", sensitivity_layer); + goto fail; + } + } + } + + orion_qwen_cpu_rmsnorm(cpu_curr + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, cpu_last); + orion_qwen_cpu_rmsnorm(hybrid_curr + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, hybrid_last); + + int cpu_top_id = -1; + int hybrid_top_id = -1; + float cpu_top_logit = -INFINITY; + float hybrid_top_logit = -INFINITY; + if (!sampled_topk_logits(blob_dir, lm_head_name, cpu_last, d_model, sample_vocab, &cpu_top_id, &cpu_top_logit) || + !sampled_topk_logits(blob_dir, lm_head_name, hybrid_last, d_model, sample_vocab, &hybrid_top_id, &hybrid_top_logit)) { + fprintf(stderr, "FAIL: topk logits failed\n"); + goto fail; + } + + double final_mean = mean_abs_diff(cpu_last, hybrid_last, d_model); + double final_max = max_abs_diff(cpu_last, hybrid_last, d_model); + + printf(" first_diverged_layer=%d\n", first_diverged_layer); + printf(" first_diverged_max_abs_diff=%.6f\n", first_diverged_max); + printf(" worst_layer=%d\n", worst_layer); + printf(" worst_layer_mean_abs_diff=%.6f\n", worst_mean); + printf(" worst_layer_max_abs_diff=%.6f\n", worst_max); + printf(" final_norm_mean_abs_diff=%.6f\n", final_mean); + printf(" final_norm_max_abs_diff=%.6f\n", final_max); + printf(" first_pair_pref_diverged_layer=%d\n", first_pair_pref_diverged_layer); + printf(" worst_pair_gap_layer=%d\n", worst_pair_gap_layer); + printf(" worst_pair_gap_abs_diff=%.6f\n", worst_pair_gap_abs_diff); + printf(" first_tail_pref_diverged_layer=%d\n", first_tail_pref_diverged_layer); + if (pair_enabled && sensitivity_layer >= 0) { + printf(" sensitivity_stage_kind=%s\n", sensitivity_stage_kind); + printf(" sensitivity_input_mean_abs_diff=%.6f\n", sensitivity_input_mean_abs_diff); + printf(" sensitivity_input_max_abs_diff=%.6f\n", sensitivity_input_max_abs_diff); + if (strcmp(sensitivity_stage_kind, "full_attention") == 0) { + printf(" cpu_input_cpu_hidden_attn_pair_gap=%.6f\n", sensitivity_cpu_input.cpu_hidden_attn_pair_gap); + printf(" cpu_input_ane_hidden_attn_pair_gap=%.6f\n", sensitivity_cpu_input.ane_hidden_attn_pair_gap); + printf(" cpu_input_cpu_final_pair_gap=%.6f\n", sensitivity_cpu_input.cpu_final_pair_gap); + printf(" cpu_input_ane_final_pair_gap=%.6f\n", sensitivity_cpu_input.ane_final_pair_gap); + printf(" cpu_input_cpu_hidden_attn_pref_token=%d\n", sensitivity_cpu_input.cpu_hidden_attn_pref_token); + printf(" cpu_input_ane_hidden_attn_pref_token=%d\n", sensitivity_cpu_input.ane_hidden_attn_pref_token); + printf(" cpu_input_cpu_final_pref_token=%d\n", sensitivity_cpu_input.cpu_final_pref_token); + printf(" cpu_input_ane_final_pref_token=%d\n", sensitivity_cpu_input.ane_final_pref_token); + printf(" hybrid_input_cpu_hidden_attn_pair_gap=%.6f\n", sensitivity_hybrid_input.cpu_hidden_attn_pair_gap); + printf(" hybrid_input_ane_hidden_attn_pair_gap=%.6f\n", sensitivity_hybrid_input.ane_hidden_attn_pair_gap); + printf(" hybrid_input_cpu_final_pair_gap=%.6f\n", sensitivity_hybrid_input.cpu_final_pair_gap); + printf(" hybrid_input_ane_final_pair_gap=%.6f\n", sensitivity_hybrid_input.ane_final_pair_gap); + printf(" hybrid_input_cpu_hidden_attn_pref_token=%d\n", sensitivity_hybrid_input.cpu_hidden_attn_pref_token); + printf(" hybrid_input_ane_hidden_attn_pref_token=%d\n", sensitivity_hybrid_input.ane_hidden_attn_pref_token); + printf(" hybrid_input_cpu_final_pref_token=%d\n", sensitivity_hybrid_input.cpu_final_pref_token); + printf(" hybrid_input_ane_final_pref_token=%d\n", sensitivity_hybrid_input.ane_final_pref_token); + printf(" full_input_normed_mean_abs_diff=%.6f\n", full_input_sensitivity.normed_mean_abs_diff); + printf(" full_input_normed_max_abs_diff=%.6f\n", full_input_sensitivity.normed_max_abs_diff); + printf(" full_input_q_proj_mean_abs_diff=%.6f\n", full_input_sensitivity.q_proj_mean_abs_diff); + printf(" full_input_q_proj_max_abs_diff=%.6f\n", full_input_sensitivity.q_proj_max_abs_diff); + printf(" full_input_q_gate_mean_abs_diff=%.6f\n", full_input_sensitivity.q_gate_mean_abs_diff); + printf(" full_input_q_gate_max_abs_diff=%.6f\n", full_input_sensitivity.q_gate_max_abs_diff); + printf(" full_input_k_proj_mean_abs_diff=%.6f\n", full_input_sensitivity.k_proj_mean_abs_diff); + printf(" full_input_k_proj_max_abs_diff=%.6f\n", full_input_sensitivity.k_proj_max_abs_diff); + printf(" full_input_v_proj_mean_abs_diff=%.6f\n", full_input_sensitivity.v_proj_mean_abs_diff); + printf(" full_input_v_proj_max_abs_diff=%.6f\n", full_input_sensitivity.v_proj_max_abs_diff); + printf(" full_input_attn_gate_sigmoid_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_gate_sigmoid_mean_abs_diff); + printf(" full_input_attn_gate_sigmoid_max_abs_diff=%.6f\n", full_input_sensitivity.attn_gate_sigmoid_max_abs_diff); + printf(" full_input_attn_context_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_context_mean_abs_diff); + printf(" full_input_attn_context_max_abs_diff=%.6f\n", full_input_sensitivity.attn_context_max_abs_diff); + printf(" full_input_attn_gated_context_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_gated_context_mean_abs_diff); + printf(" full_input_attn_gated_context_max_abs_diff=%.6f\n", full_input_sensitivity.attn_gated_context_max_abs_diff); + printf(" full_input_attn_out_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_out_mean_abs_diff); + printf(" full_input_attn_out_max_abs_diff=%.6f\n", full_input_sensitivity.attn_out_max_abs_diff); + printf(" full_input_attn_qgate_only_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_qgate_only_mean_abs_diff); + printf(" full_input_attn_qgate_only_max_abs_diff=%.6f\n", full_input_sensitivity.attn_qgate_only_max_abs_diff); + printf(" full_input_attn_v_only_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_v_only_mean_abs_diff); + printf(" full_input_attn_v_only_max_abs_diff=%.6f\n", full_input_sensitivity.attn_v_only_max_abs_diff); + printf(" full_input_attn_context_only_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_context_only_mean_abs_diff); + printf(" full_input_attn_context_only_max_abs_diff=%.6f\n", full_input_sensitivity.attn_context_only_max_abs_diff); + printf(" full_input_attn_sigmoid_only_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_sigmoid_only_mean_abs_diff); + printf(" full_input_attn_sigmoid_only_max_abs_diff=%.6f\n", full_input_sensitivity.attn_sigmoid_only_max_abs_diff); + printf(" full_input_attn_gated_context_only_mean_abs_diff=%.6f\n", full_input_sensitivity.attn_gated_context_only_mean_abs_diff); + printf(" full_input_attn_gated_context_only_max_abs_diff=%.6f\n", full_input_sensitivity.attn_gated_context_only_max_abs_diff); + printf(" full_input_hidden_attn_mean_abs_diff=%.6f\n", full_input_sensitivity.hidden_attn_mean_abs_diff); + printf(" full_input_hidden_attn_max_abs_diff=%.6f\n", full_input_sensitivity.hidden_attn_max_abs_diff); + printf(" full_input_ffn_rms_mean_abs_diff=%.6f\n", full_input_sensitivity.ffn_rms_mean_abs_diff); + printf(" full_input_ffn_rms_max_abs_diff=%.6f\n", full_input_sensitivity.ffn_rms_max_abs_diff); + printf(" full_input_ffn_gate_mean_abs_diff=%.6f\n", full_input_sensitivity.ffn_gate_mean_abs_diff); + printf(" full_input_ffn_gate_max_abs_diff=%.6f\n", full_input_sensitivity.ffn_gate_max_abs_diff); + printf(" full_input_ffn_up_mean_abs_diff=%.6f\n", full_input_sensitivity.ffn_up_mean_abs_diff); + printf(" full_input_ffn_up_max_abs_diff=%.6f\n", full_input_sensitivity.ffn_up_max_abs_diff); + printf(" full_input_ffn_silu_mean_abs_diff=%.6f\n", full_input_sensitivity.ffn_silu_mean_abs_diff); + printf(" full_input_ffn_silu_max_abs_diff=%.6f\n", full_input_sensitivity.ffn_silu_max_abs_diff); + printf(" full_input_ffn_hidden_mean_abs_diff=%.6f\n", full_input_sensitivity.ffn_hidden_mean_abs_diff); + printf(" full_input_ffn_hidden_max_abs_diff=%.6f\n", full_input_sensitivity.ffn_hidden_max_abs_diff); + printf(" full_input_ffn_down_mean_abs_diff=%.6f\n", full_input_sensitivity.ffn_down_mean_abs_diff); + printf(" full_input_ffn_down_max_abs_diff=%.6f\n", full_input_sensitivity.ffn_down_max_abs_diff); + printf(" full_input_final_mean_abs_diff=%.6f\n", full_input_sensitivity.final_mean_abs_diff); + printf(" full_input_final_max_abs_diff=%.6f\n", full_input_sensitivity.final_max_abs_diff); + printf(" full_input_cpu_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.cpu_hidden_attn_pair_gap); + printf(" full_input_hybrid_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.hybrid_hidden_attn_pair_gap); + printf(" full_input_cpu_res_hybrid_qgate_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_qgate_hidden_attn_pair_gap); + printf(" full_input_cpu_res_hybrid_v_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_v_hidden_attn_pair_gap); + printf(" full_input_cpu_res_hybrid_context_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_context_hidden_attn_pair_gap); + printf(" full_input_cpu_res_hybrid_sigmoid_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_sigmoid_hidden_attn_pair_gap); + printf(" full_input_cpu_res_hybrid_gated_context_hidden_attn_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_gated_context_hidden_attn_pair_gap); + printf(" full_input_cpu_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_final_pair_gap); + printf(" full_input_hybrid_final_pair_gap=%.6f\n", full_input_sensitivity.hybrid_final_pair_gap); + printf(" full_input_cpu_res_hybrid_qgate_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_qgate_final_pair_gap); + printf(" full_input_cpu_res_hybrid_v_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_v_final_pair_gap); + printf(" full_input_cpu_res_hybrid_context_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_context_final_pair_gap); + printf(" full_input_cpu_res_hybrid_sigmoid_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_sigmoid_final_pair_gap); + printf(" full_input_cpu_res_hybrid_gated_context_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_gated_context_final_pair_gap); + printf(" full_input_cpu_hidden_attn_pref_token=%d\n", full_input_sensitivity.cpu_hidden_attn_pref_token); + printf(" full_input_hybrid_hidden_attn_pref_token=%d\n", full_input_sensitivity.hybrid_hidden_attn_pref_token); + printf(" full_input_cpu_res_hybrid_qgate_hidden_attn_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_qgate_hidden_attn_pref_token); + printf(" full_input_cpu_res_hybrid_v_hidden_attn_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_v_hidden_attn_pref_token); + printf(" full_input_cpu_res_hybrid_context_hidden_attn_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_context_hidden_attn_pref_token); + printf(" full_input_cpu_res_hybrid_sigmoid_hidden_attn_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_sigmoid_hidden_attn_pref_token); + printf(" full_input_cpu_res_hybrid_gated_context_hidden_attn_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_gated_context_hidden_attn_pref_token); + printf(" full_input_cpu_final_pref_token=%d\n", full_input_sensitivity.cpu_final_pref_token); + printf(" full_input_hybrid_final_pref_token=%d\n", full_input_sensitivity.hybrid_final_pref_token); + printf(" full_input_cpu_res_hybrid_qgate_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_qgate_final_pref_token); + printf(" full_input_cpu_res_hybrid_v_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_v_final_pref_token); + printf(" full_input_cpu_res_hybrid_context_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_context_final_pref_token); + printf(" full_input_cpu_res_hybrid_sigmoid_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_sigmoid_final_pref_token); + printf(" full_input_cpu_res_hybrid_gated_context_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_gated_context_final_pref_token); + printf(" full_input_attn_v_attr_count=%d\n", full_input_sensitivity.attn_v_attr_count); + for (int rank = 0; rank < full_input_sensitivity.attn_v_attr_count; rank++) { + printf(" full_input_attn_v_attr_rank=%d dim=%d cpu_norm=%.6f hybrid_norm=%.6f delta_norm=%.6f pair_delta=%.6f abs_contrib=%.6f signed_contrib=%.6f\n", + rank + 1, + full_input_sensitivity.attn_v_attr_dims[rank], + full_input_sensitivity.attn_v_attr_cpu_norm[rank], + full_input_sensitivity.attn_v_attr_hybrid_norm[rank], + full_input_sensitivity.attn_v_attr_delta_norm[rank], + full_input_sensitivity.attn_v_attr_pair_delta[rank], + full_input_sensitivity.attn_v_attr_abs_contrib[rank], + full_input_sensitivity.attn_v_attr_signed_contrib[rank]); + } + printf(" full_input_attn_gated_context_attr_count=%d\n", full_input_sensitivity.attn_gated_context_attr_count); + for (int rank = 0; rank < full_input_sensitivity.attn_gated_context_attr_count; rank++) { + printf(" full_input_attn_gated_context_attr_rank=%d dim=%d cpu_norm=%.6f hybrid_norm=%.6f delta_norm=%.6f pair_delta=%.6f abs_contrib=%.6f signed_contrib=%.6f\n", + rank + 1, + full_input_sensitivity.attn_gated_context_attr_dims[rank], + full_input_sensitivity.attn_gated_context_attr_cpu_norm[rank], + full_input_sensitivity.attn_gated_context_attr_hybrid_norm[rank], + full_input_sensitivity.attn_gated_context_attr_delta_norm[rank], + full_input_sensitivity.attn_gated_context_attr_pair_delta[rank], + full_input_sensitivity.attn_gated_context_attr_abs_contrib[rank], + full_input_sensitivity.attn_gated_context_attr_signed_contrib[rank]); + } + printf(" full_input_ffn_down_attr_count=%d\n", full_input_sensitivity.ffn_down_attr_count); + for (int rank = 0; rank < full_input_sensitivity.ffn_down_attr_count; rank++) { + printf(" full_input_ffn_down_attr_rank=%d dim=%d cpu_norm=%.6f hybrid_norm=%.6f delta_norm=%.6f pair_delta=%.6f abs_contrib=%.6f signed_contrib=%.6f\n", + rank + 1, + full_input_sensitivity.ffn_down_attr_dims[rank], + full_input_sensitivity.ffn_down_attr_cpu_norm[rank], + full_input_sensitivity.ffn_down_attr_hybrid_norm[rank], + full_input_sensitivity.ffn_down_attr_delta_norm[rank], + full_input_sensitivity.ffn_down_attr_pair_delta[rank], + full_input_sensitivity.ffn_down_attr_abs_contrib[rank], + full_input_sensitivity.ffn_down_attr_signed_contrib[rank]); + } + printf(" full_input_cpu_res_hybrid_down_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_down_final_pair_gap); + printf(" full_input_hybrid_res_cpu_down_final_pair_gap=%.6f\n", full_input_sensitivity.hybrid_res_cpu_down_final_pair_gap); + printf(" full_input_cpu_res_hybrid_down_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_down_final_pref_token); + printf(" full_input_hybrid_res_cpu_down_final_pref_token=%d\n", full_input_sensitivity.hybrid_res_cpu_down_final_pref_token); + printf(" full_input_cpu_res_cpu_silu_hybrid_up_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_cpu_silu_hybrid_up_final_pair_gap); + printf(" full_input_cpu_res_hybrid_silu_cpu_up_final_pair_gap=%.6f\n", full_input_sensitivity.cpu_res_hybrid_silu_cpu_up_final_pair_gap); + printf(" full_input_cpu_res_cpu_silu_hybrid_up_final_pref_token=%d\n", full_input_sensitivity.cpu_res_cpu_silu_hybrid_up_final_pref_token); + printf(" full_input_cpu_res_hybrid_silu_cpu_up_final_pref_token=%d\n", full_input_sensitivity.cpu_res_hybrid_silu_cpu_up_final_pref_token); + printf(" full_input_gate_attr_total_channels=%d\n", full_input_sensitivity.gate_attr_total_channels); + printf(" full_input_gate_attr_report_count=%d\n", full_input_sensitivity.gate_attr_report_count); + printf(" full_input_gate_attr_first_cumulative_flip_rank=%d\n", full_input_sensitivity.gate_attr_first_cumulative_flip_rank); + printf(" full_input_gate_attr_flip_channel=%d\n", full_input_sensitivity.gate_attr_flip_channel); + printf(" full_input_gate_attr_flip_delta_hidden=%.6f\n", full_input_sensitivity.gate_attr_flip_delta_hidden); + printf(" full_input_gate_attr_flip_cumulative_pair_gap=%.6f\n", full_input_sensitivity.gate_attr_flip_cumulative_pair_gap); + printf(" full_input_gate_attr_all_cumulative_pair_gap=%.6f\n", full_input_sensitivity.gate_attr_all_cumulative_pair_gap); + printf(" full_input_gate_input_attr_channel_limit=%d\n", full_input_sensitivity.gate_input_attr_channel_limit); + printf(" full_input_gate_input_attr_count=%d\n", full_input_sensitivity.gate_input_attr_count); + for (int rank = 0; rank < full_input_sensitivity.gate_attr_report_count; rank++) { + printf(" full_input_gate_attr_rank=%d channel=%d cpu_gate=%.6f hybrid_gate=%.6f delta_gate=%.6f cpu_silu=%.6f hybrid_silu=%.6f delta_silu=%.6f cpu_up=%.6f delta_hidden=%.6f single_pair_gap=%.6f single_pref_token=%d cumulative_pair_gap=%.6f cumulative_pref_token=%d\n", + rank + 1, + full_input_sensitivity.gate_attr_channels[rank], + full_input_sensitivity.gate_attr_cpu_gate[rank], + full_input_sensitivity.gate_attr_hybrid_gate[rank], + full_input_sensitivity.gate_attr_delta_gate[rank], + full_input_sensitivity.gate_attr_cpu_silu[rank], + full_input_sensitivity.gate_attr_hybrid_silu[rank], + full_input_sensitivity.gate_attr_delta_silu[rank], + full_input_sensitivity.gate_attr_cpu_up[rank], + full_input_sensitivity.gate_attr_delta_hidden[rank], + full_input_sensitivity.gate_attr_single_pair_gap[rank], + full_input_sensitivity.gate_attr_single_pref_token[rank], + full_input_sensitivity.gate_attr_cumulative_pair_gap[rank], + full_input_sensitivity.gate_attr_cumulative_pref_token[rank]); + } + for (int rank = 0; rank < full_input_sensitivity.gate_input_attr_count; rank++) { + printf(" full_input_gate_input_attr_rank=%d dim=%d delta_ffn_rms=%.6f abs_contrib=%.6f signed_contrib=%.6f\n", + rank + 1, + full_input_sensitivity.gate_input_attr_dims[rank], + full_input_sensitivity.gate_input_attr_delta_ffn_rms[rank], + full_input_sensitivity.gate_input_attr_abs_contrib[rank], + full_input_sensitivity.gate_input_attr_signed_contrib[rank]); + } + } else if (strcmp(sensitivity_stage_kind, "linear_attention") == 0) { + printf(" linear_normed_mean_abs_diff=%.6f\n", linear_sensitivity.normed_mean_abs_diff); + printf(" linear_normed_max_abs_diff=%.6f\n", linear_sensitivity.normed_max_abs_diff); + printf(" linear_mixed_linear_mean_abs_diff=%.6f\n", linear_sensitivity.mixed_linear_mean_abs_diff); + printf(" linear_mixed_linear_max_abs_diff=%.6f\n", linear_sensitivity.mixed_linear_max_abs_diff); + printf(" linear_mixed_conv_mean_abs_diff=%.6f\n", linear_sensitivity.mixed_conv_mean_abs_diff); + printf(" linear_mixed_conv_max_abs_diff=%.6f\n", linear_sensitivity.mixed_conv_max_abs_diff); + printf(" linear_query_mean_abs_diff=%.6f\n", linear_sensitivity.query_mean_abs_diff); + printf(" linear_query_max_abs_diff=%.6f\n", linear_sensitivity.query_max_abs_diff); + printf(" linear_key_mean_abs_diff=%.6f\n", linear_sensitivity.key_mean_abs_diff); + printf(" linear_key_max_abs_diff=%.6f\n", linear_sensitivity.key_max_abs_diff); + printf(" linear_value_mean_abs_diff=%.6f\n", linear_sensitivity.value_mean_abs_diff); + printf(" linear_value_max_abs_diff=%.6f\n", linear_sensitivity.value_max_abs_diff); + printf(" linear_z_mean_abs_diff=%.6f\n", linear_sensitivity.z_mean_abs_diff); + printf(" linear_z_max_abs_diff=%.6f\n", linear_sensitivity.z_max_abs_diff); + printf(" linear_beta_mean_abs_diff=%.6f\n", linear_sensitivity.beta_mean_abs_diff); + printf(" linear_beta_max_abs_diff=%.6f\n", linear_sensitivity.beta_max_abs_diff); + printf(" linear_g_mean_abs_diff=%.6f\n", linear_sensitivity.g_mean_abs_diff); + printf(" linear_g_max_abs_diff=%.6f\n", linear_sensitivity.g_max_abs_diff); + printf(" linear_core_pre_mean_abs_diff=%.6f\n", linear_sensitivity.core_pre_mean_abs_diff); + printf(" linear_core_pre_max_abs_diff=%.6f\n", linear_sensitivity.core_pre_max_abs_diff); + printf(" linear_core_mean_abs_diff=%.6f\n", linear_sensitivity.core_mean_abs_diff); + printf(" linear_core_max_abs_diff=%.6f\n", linear_sensitivity.core_max_abs_diff); + printf(" linear_attn_out_mean_abs_diff=%.6f\n", linear_sensitivity.attn_out_mean_abs_diff); + printf(" linear_attn_out_max_abs_diff=%.6f\n", linear_sensitivity.attn_out_max_abs_diff); + printf(" linear_hidden_attn_mean_abs_diff=%.6f\n", linear_sensitivity.hidden_attn_mean_abs_diff); + printf(" linear_hidden_attn_max_abs_diff=%.6f\n", linear_sensitivity.hidden_attn_max_abs_diff); + printf(" linear_final_mean_abs_diff=%.6f\n", linear_sensitivity.final_mean_abs_diff); + printf(" linear_final_max_abs_diff=%.6f\n", linear_sensitivity.final_max_abs_diff); + printf(" linear_cpu_hidden_attn_pair_gap=%.6f\n", linear_sensitivity.cpu_hidden_attn_pair_gap); + printf(" linear_hybrid_hidden_attn_pair_gap=%.6f\n", linear_sensitivity.hybrid_hidden_attn_pair_gap); + printf(" linear_cpu_final_pair_gap=%.6f\n", linear_sensitivity.cpu_final_pair_gap); + printf(" linear_hybrid_final_pair_gap=%.6f\n", linear_sensitivity.hybrid_final_pair_gap); + printf(" linear_cpu_hidden_attn_pref_token=%d\n", linear_sensitivity.cpu_hidden_attn_pref_token); + printf(" linear_hybrid_hidden_attn_pref_token=%d\n", linear_sensitivity.hybrid_hidden_attn_pref_token); + printf(" linear_cpu_final_pref_token=%d\n", linear_sensitivity.cpu_final_pref_token); + printf(" linear_hybrid_final_pref_token=%d\n", linear_sensitivity.hybrid_final_pref_token); + } + } + printf(" cpu_top_id=%d\n", cpu_top_id); + printf(" cpu_top_logit=%.6f\n", cpu_top_logit); + printf(" hybrid_top_id=%d\n", hybrid_top_id); + printf(" hybrid_top_logit=%.6f\n", hybrid_top_logit); + printf("PASS: qwen35 9b hybrid layer diff trace\n"); + + free(final_norm); + free(cpu_curr); free(cpu_next); free(hybrid_curr); free(hybrid_next); free(sensitivity_cpu_in); free(sensitivity_hybrid_in); free(cpu_last); free(hybrid_last); + if (sensitivity_bridge_owned) bridge_release(&sensitivity_bridge); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 0; + +fail: + free(final_norm); + free(cpu_curr); free(cpu_next); free(hybrid_curr); free(hybrid_next); free(sensitivity_cpu_in); free(sensitivity_hybrid_in); free(cpu_last); free(hybrid_last); + if (sensitivity_bridge_owned) bridge_release(&sensitivity_bridge); + for (int i = 0; i < manifest->n_layer; i++) bridge_release(&bridges[i]); + free(bridges); + free(bridge_mask); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } +} diff --git a/tests/test_qwen35_9b_lora_adapter_roundtrip.m b/tests/test_qwen35_9b_lora_adapter_roundtrip.m new file mode 100644 index 0000000..6125526 --- /dev/null +++ b/tests/test_qwen35_9b_lora_adapter_roundtrip.m @@ -0,0 +1,57 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + const char *adapter_dir = argv[2]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + OrionQwen9BLoRATrainer lhs; + OrionQwen9BLoRATrainer rhs; + orion_qwen9b_lora_trainer_init(&lhs, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 20260314u); + orion_qwen9b_lora_trainer_init(&rhs, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1u); + lhs.step = 3; + if (!orion_qwen9b_lora_trainer_save(&lhs, adapter_dir)) { + fprintf(stderr, "FAIL: save failed\n"); + goto fail; + } + if (!orion_qwen9b_lora_trainer_load(&rhs, adapter_dir)) { + fprintf(stderr, "FAIL: load failed\n"); + goto fail; + } + if (!orion_qwen9b_lora_trainer_compare(&lhs, &rhs, 0.0f)) { + fprintf(stderr, "FAIL: roundtrip mismatch\n"); + goto fail; + } + + printf("PASS: qwen35 9b lora adapter roundtrip\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" adapter_dir=%s\n", adapter_dir); + printf(" rank=%d\n", lhs.q_proj.rank); + printf(" step=%d\n", rhs.step); + + orion_qwen9b_lora_trainer_free(&lhs); + orion_qwen9b_lora_trainer_free(&rhs); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + orion_qwen9b_lora_trainer_free(&lhs); + orion_qwen9b_lora_trainer_free(&rhs); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_9b_lora_ane_forward_probe.m b/tests/test_qwen35_9b_lora_ane_forward_probe.m new file mode 100644 index 0000000..088eb50 --- /dev/null +++ b/tests/test_qwen35_9b_lora_ane_forward_probe.m @@ -0,0 +1,90 @@ +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static float abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return (float)total; +} + +static float max_abs(const float *x, int n) { + float best = 0.0f; + for (int i = 0; i < n; i++) { + float v = fabsf(x[i]); + if (v > best) best = v; + } + return best; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s [input_token]\n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + const int input_token = argc >= 3 ? atoi(argv[2]) : 0; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + const int layer_idx = manifest->n_layer - 1; + const int d_model = manifest->d_model; + float *hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *input_ln = NULL; + if (!hidden || !normed) { + fprintf(stderr, "FAIL: alloc failed\n"); + free(hidden); + free(normed); + orion_qwen35_manifest_free(manifest); + return 1; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/input_layernorm.bin", blob_dir, layer_idx); + input_ln = orion_read_blob_f32_exact(path, d_model); + if (!input_ln) { + fprintf(stderr, "FAIL: input layernorm load failed\n"); + free(hidden); + free(normed); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, input_token, hidden)) { + fprintf(stderr, "FAIL: frozen prefix hidden probe failed\n"); + free(input_ln); + free(hidden); + free(normed); + orion_qwen35_manifest_free(manifest); + return 1; + } + + orion_qwen_cpu_rmsnorm(hidden, input_ln, d_model, 1e-6f, normed); + + printf("PASS: qwen35 9b lora ane forward probe\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" input_token=%d\n", input_token); + printf(" layer_idx=%d\n", layer_idx); + printf(" hidden_abs_sum=%.6f\n", abs_sum(hidden, d_model)); + printf(" hidden_max_abs=%.6f\n", max_abs(hidden, d_model)); + printf(" normed_abs_sum=%.6f\n", abs_sum(normed, d_model)); + printf(" normed_max_abs=%.6f\n", max_abs(normed, d_model)); + printf(" next_blocker=%s\n", "wire ANE q_proj/kv_proj forward on top of frozen prefix hidden"); + + free(input_ln); + free(hidden); + free(normed); + orion_qwen35_manifest_free(manifest); + return 0; + } +} diff --git a/tests/test_qwen35_9b_lora_ane_qv_probe.m b/tests/test_qwen35_9b_lora_ane_qv_probe.m new file mode 100644 index 0000000..5964a87 --- /dev/null +++ b/tests/test_qwen35_9b_lora_ane_qv_probe.m @@ -0,0 +1,477 @@ +#import +#import +#import +#include +#include +#include +#include +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" +#import "../kernels/training/qwen_lora_train.h" + +static float abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return (float)total; +} + +static float max_abs(const float *x, int n) { + float best = 0.0f; + for (int i = 0; i < n; i++) { + float v = fabsf(x[i]); + if (v > best) best = v; + } + return best; +} + +static double mean_abs_diff(const float *a, const float *b, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)a[i] - (double)b[i]); + return total / (double)n; +} + +static double max_abs_diff(const float *a, const float *b, int n) { + double best = 0.0; + for (int i = 0; i < n; i++) { + double d = fabs((double)a[i] - (double)b[i]); + if (d > best) best = d; + } + return best; +} + +static double mean_abs_diff_q_half(const float *a, const float *b, int q_dim, int half_idx) { + const float *row_a = a + half_idx * q_dim; + const float *row_b = b + half_idx * q_dim; + return mean_abs_diff(row_a, row_b, q_dim); +} + +static double max_abs_diff_q_half(const float *a, const float *b, int q_dim, int half_idx) { + const float *row_a = a + half_idx * q_dim; + const float *row_b = b + half_idx * q_dim; + return max_abs_diff(row_a, row_b, q_dim); +} + +static double mean_abs_diff_q_half_swapped(const float *cpu_q, const float *ane_q, int q_dim) { + double total = 0.0; + const float *cpu_query = cpu_q; + const float *cpu_gate = cpu_q + q_dim; + const float *ane_gate = ane_q; + const float *ane_query = ane_q + q_dim; + for (int i = 0; i < q_dim; i++) total += fabs((double)cpu_query[i] - (double)ane_query[i]); + for (int i = 0; i < q_dim; i++) total += fabs((double)cpu_gate[i] - (double)ane_gate[i]); + return total / (double)(q_dim * 2); +} + +static double max_abs_diff_q_half_swapped(const float *cpu_q, const float *ane_q, int q_dim) { + double best = 0.0; + const float *cpu_query = cpu_q; + const float *cpu_gate = cpu_q + q_dim; + const float *ane_gate = ane_q; + const float *ane_query = ane_q + q_dim; + for (int i = 0; i < q_dim; i++) { + double d = fabs((double)cpu_query[i] - (double)ane_query[i]); + if (d > best) best = d; + } + for (int i = 0; i < q_dim; i++) { + double d = fabs((double)cpu_gate[i] - (double)ane_gate[i]); + if (d > best) best = d; + } + return best; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static void linear_seq(const float *x_seq, + int seq_len, + int in_dim, + const float *weight, + int out_dim, + float *out_seq) { + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + seq_len, out_dim, in_dim, + 1.0f, x_seq, in_dim, weight, in_dim, + 0.0f, out_seq, out_dim); +} + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_pattern_input_surface(int d_model, int bucket) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < bucket; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = sinf((float)(c * 0.001 + t * 0.01)) * 0.1f; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_qproj_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_linear_only_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s [input_token] [layer_idx] [bucket]\n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + const int input_token = argc >= 3 ? atoi(argv[2]) : 0; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + const int layer_idx = (argc >= 4 && argv[3][0] != '\0') ? atoi(argv[3]) : (manifest->n_layer - 1); + const int seq_len = 1; + const int bucket = argc >= 5 ? atoi(argv[4]) : 32; + const int d_model = manifest->d_model; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head > 0 ? manifest->n_kv_head : manifest->n_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + NSString *blobDir = [NSString stringWithUTF8String:blob_dir]; + + float *hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *cpu_q = (float *)calloc((size_t)q_dim * 2, sizeof(float)); + float *cpu_k = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *cpu_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *ane_q = (float *)calloc((size_t)q_dim * 2, sizeof(float)); + float *ane_k = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *ane_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *cpu_attn = (float *)calloc((size_t)d_model, sizeof(float)); + float *ane_attn = (float *)calloc((size_t)d_model, sizeof(float)); + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", q_dim * 2 * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + IOSurfaceRef ioIn = NULL; + IOSurfaceRef ioNormedIn = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + OrionProgram *progQ = NULL; + OrionProgram *progKV = NULL; + OrionProgram *progQLinear = NULL; + OrionProgram *progKVLinear = NULL; + NSString *milQ = nil; + NSString *milKV = nil; + NSString *milQLinear = nil; + NSString *milKVLinear = nil; + IOSurfaceRef ioPatternIn = NULL; + const char *graph_mode = "raw_hidden_full"; + int raw_hidden_eval_ok = 0; + int raw_hidden_q_eval_ok = 0; + int raw_hidden_kv_eval_ok = 0; + int cpu_rmsnorm_q_eval_ok = 0; + int cpu_rmsnorm_kv_eval_ok = 0; + int pattern_q_eval_ok = 0; + int pattern_q_linear_eval_ok = 0; + int exit_code = 1; + + if (!hidden || !normed || !cpu_q || !cpu_k || !cpu_v || !ane_q || !ane_k || !ane_v || + !cpu_attn || !ane_attn || !input_ln || !q_proj || !k_proj || !v_proj || !o_proj || + !q_norm || !k_norm) { + fprintf(stderr, "FAIL: alloc or weight load failed\n"); + goto cleanup; + } + + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, input_token, hidden)) { + fprintf(stderr, "FAIL: frozen prefix hidden probe failed\n"); + goto cleanup; + } + + orion_qwen_cpu_rmsnorm(hidden, input_ln, d_model, 1e-6f, normed); + linear_seq(normed, seq_len, d_model, q_proj, q_dim * 2, cpu_q); + linear_seq(normed, seq_len, d_model, k_proj, kv_dim, cpu_k); + linear_seq(normed, seq_len, d_model, v_proj, kv_dim, cpu_v); + orion_qwen_cpu_full_attention_from_projections_with_rope( + cpu_q, cpu_k, cpu_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + cpu_attn + ); + + if (!orion_ane_init()) { + fprintf(stderr, "FAIL: ane init failed\n"); + goto cleanup; + } + + const int compile_before = orion_compile_count(); + OrionModelConfig cfg = { + .n_layer = manifest->n_layer, + .n_head = manifest->n_head, + .n_kv_head = n_kv_head, + .d_model = manifest->d_model, + .head_dim = manifest->head_dim, + .hidden_dim = manifest->d_ff, + .vocab = manifest->vocab, + .max_seq = manifest->max_seq, + }; + milQ = compile_graph(orion_frontend_qwen35_prefill_q_proj(layer_idx, bucket, &cfg)); + milKV = compile_graph(orion_frontend_qwen35_prefill_kv_proj(layer_idx, bucket, &cfg)); + milQLinear = compile_graph(orion_frontend_qwen35_prefill_q_proj_linear_only(layer_idx, bucket, &cfg)); + milKVLinear = compile_graph(orion_frontend_qwen35_prefill_kv_proj_linear_only(layer_idx, bucket, &cfg)); + if (!milQ || !milKV || !milQLinear || !milKVLinear) { + fprintf(stderr, "FAIL: mil graph build failed\n"); + goto cleanup; + } + + progQ = orion_compile_mil(milQ.UTF8String, build_qproj_wdict(layer_idx, blobDir), "qwen35_9b_lora_probe_q"); + progKV = orion_compile_mil(milKV.UTF8String, build_kv_wdict(layer_idx, blobDir), "qwen35_9b_lora_probe_kv"); + progQLinear = orion_compile_mil(milQLinear.UTF8String, build_qproj_linear_only_wdict(layer_idx, blobDir), "qwen35_9b_lora_probe_q_linear"); + progKVLinear = orion_compile_mil(milKVLinear.UTF8String, build_kv_linear_only_wdict(layer_idx, blobDir), "qwen35_9b_lora_probe_kv_linear"); + if (!progQ || !progKV || !progQLinear || !progKVLinear) { + fprintf(stderr, "FAIL: ane compile failed\n"); + goto cleanup; + } + const int compile_after = orion_compile_count(); + + ioIn = make_cpu_seq_input_surface(hidden, seq_len, bucket, d_model); + ioNormedIn = make_cpu_seq_input_surface(normed, seq_len, bucket, d_model); + ioPatternIn = make_pattern_input_surface(d_model, bucket); + ioQ = make_f32_surface(q_dim * 2 * bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bucket, 0.0f); + if (!ioIn || !ioNormedIn || !ioPatternIn || !ioQ || !ioK || !ioV) { + fprintf(stderr, "FAIL: iosurface alloc failed\n"); + goto cleanup; + } + + IOSurfaceRef insQ[] = { ioIn }; + IOSurfaceRef insQLinear[] = { ioNormedIn }; + IOSurfaceRef insQPattern[] = { ioPatternIn }; + IOSurfaceRef outsQ[] = { ioQ }; + IOSurfaceRef outsKV[] = { ioK, ioV }; + raw_hidden_q_eval_ok = orion_eval(progQ, insQ, 1, outsQ, 1) ? 1 : 0; + raw_hidden_kv_eval_ok = orion_eval(progKV, insQ, 1, outsKV, 2) ? 1 : 0; + if (raw_hidden_q_eval_ok && raw_hidden_kv_eval_ok) { + raw_hidden_eval_ok = 1; + graph_mode = "raw_hidden_full"; + } else { + graph_mode = "cpu_rmsnorm_linear_only"; + cpu_rmsnorm_q_eval_ok = orion_eval(progQLinear, insQLinear, 1, outsQ, 1) ? 1 : 0; + cpu_rmsnorm_kv_eval_ok = orion_eval(progKVLinear, insQLinear, 1, outsKV, 2) ? 1 : 0; + if (!(cpu_rmsnorm_q_eval_ok && cpu_rmsnorm_kv_eval_ok)) { + pattern_q_eval_ok = orion_eval(progQ, insQPattern, 1, outsQ, 1) ? 1 : 0; + pattern_q_linear_eval_ok = orion_eval(progQLinear, insQPattern, 1, outsQ, 1) ? 1 : 0; + printf("FAIL: qwen35 9b lora ane qv probe\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" input_token=%d\n", input_token); + printf(" layer_idx=%d\n", layer_idx); + printf(" bucket=%d\n", bucket); + printf(" graph_mode=%s\n", graph_mode); + printf(" raw_hidden_eval_ok=%d\n", raw_hidden_eval_ok); + printf(" raw_hidden_q_eval_ok=%d\n", raw_hidden_q_eval_ok); + printf(" raw_hidden_kv_eval_ok=%d\n", raw_hidden_kv_eval_ok); + printf(" cpu_rmsnorm_q_eval_ok=%d\n", cpu_rmsnorm_q_eval_ok); + printf(" cpu_rmsnorm_kv_eval_ok=%d\n", cpu_rmsnorm_kv_eval_ok); + printf(" pattern_q_eval_ok=%d\n", pattern_q_eval_ok); + printf(" pattern_q_linear_eval_ok=%d\n", pattern_q_linear_eval_ok); + printf(" hidden_abs_sum=%.6f\n", abs_sum(hidden, d_model)); + printf(" hidden_max_abs=%.6f\n", max_abs(hidden, d_model)); + printf(" normed_abs_sum=%.6f\n", abs_sum(normed, d_model)); + printf(" normed_max_abs=%.6f\n", max_abs(normed, d_model)); + printf(" compile_count_before=%d\n", compile_before); + printf(" compile_count_after=%d\n", compile_after); + printf(" compile_count_delta=%d\n", compile_after - compile_before); + printf(" next_blocker=%s\n", "debug layer31 q graph eval path before ANE train splice"); + goto cleanup; + } + } + + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bucket, ane_q); + read_ane_surface_prefix(ioK, kv_dim, seq_len, bucket, ane_k); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bucket, ane_v); + orion_qwen_cpu_full_attention_from_projections_with_rope( + ane_q, ane_k, ane_v, seq_len, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, + manifest->rope_theta, manifest->partial_rotary_factor, + ane_attn + ); + + printf("PASS: qwen35 9b lora ane qv probe\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" input_token=%d\n", input_token); + printf(" layer_idx=%d\n", layer_idx); + printf(" bucket=%d\n", bucket); + printf(" graph_mode=%s\n", graph_mode); + printf(" raw_hidden_eval_ok=%d\n", raw_hidden_eval_ok); + printf(" raw_hidden_q_eval_ok=%d\n", raw_hidden_q_eval_ok); + printf(" raw_hidden_kv_eval_ok=%d\n", raw_hidden_kv_eval_ok); + printf(" cpu_rmsnorm_q_eval_ok=%d\n", cpu_rmsnorm_q_eval_ok); + printf(" cpu_rmsnorm_kv_eval_ok=%d\n", cpu_rmsnorm_kv_eval_ok); + printf(" pattern_q_eval_ok=%d\n", pattern_q_eval_ok); + printf(" pattern_q_linear_eval_ok=%d\n", pattern_q_linear_eval_ok); + printf(" hidden_abs_sum=%.6f\n", abs_sum(hidden, d_model)); + printf(" hidden_max_abs=%.6f\n", max_abs(hidden, d_model)); + printf(" normed_abs_sum=%.6f\n", abs_sum(normed, d_model)); + printf(" normed_max_abs=%.6f\n", max_abs(normed, d_model)); + printf(" q_proj_mean_abs_diff=%.6f\n", mean_abs_diff(cpu_q, ane_q, q_dim * 2)); + printf(" q_proj_max_abs_diff=%.6f\n", max_abs_diff(cpu_q, ane_q, q_dim * 2)); + printf(" q_query_mean_abs_diff=%.6f\n", mean_abs_diff_q_half(cpu_q, ane_q, q_dim, 0)); + printf(" q_query_max_abs_diff=%.6f\n", max_abs_diff_q_half(cpu_q, ane_q, q_dim, 0)); + printf(" q_gate_mean_abs_diff=%.6f\n", mean_abs_diff_q_half(cpu_q, ane_q, q_dim, 1)); + printf(" q_gate_max_abs_diff=%.6f\n", max_abs_diff_q_half(cpu_q, ane_q, q_dim, 1)); + printf(" q_half_swap_mean_abs_diff=%.6f\n", mean_abs_diff_q_half_swapped(cpu_q, ane_q, q_dim)); + printf(" q_half_swap_max_abs_diff=%.6f\n", max_abs_diff_q_half_swapped(cpu_q, ane_q, q_dim)); + printf(" k_proj_mean_abs_diff=%.6f\n", mean_abs_diff(cpu_k, ane_k, kv_dim)); + printf(" k_proj_max_abs_diff=%.6f\n", max_abs_diff(cpu_k, ane_k, kv_dim)); + printf(" v_proj_mean_abs_diff=%.6f\n", mean_abs_diff(cpu_v, ane_v, kv_dim)); + printf(" v_proj_max_abs_diff=%.6f\n", max_abs_diff(cpu_v, ane_v, kv_dim)); + printf(" attn_out_mean_abs_diff=%.6f\n", mean_abs_diff(cpu_attn, ane_attn, d_model)); + printf(" attn_out_max_abs_diff=%.6f\n", max_abs_diff(cpu_attn, ane_attn, d_model)); + printf(" compile_count_before=%d\n", compile_before); + printf(" compile_count_after=%d\n", compile_after); + printf(" compile_count_delta=%d\n", compile_after - compile_before); + printf(" next_blocker=%s\n", + raw_hidden_eval_ok + ? "wire CPU optimizer on top of ANE q/v forward" + : "wire CPU RMSNorm plus ANE q/v forward and debug raw hidden path"); + exit_code = 0; + +cleanup: + if (ioIn) CFRelease(ioIn); + if (ioNormedIn) CFRelease(ioNormedIn); + if (ioPatternIn) CFRelease(ioPatternIn); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (progQ) orion_release_program(progQ); + if (progKV) orion_release_program(progKV); + if (progQLinear) orion_release_program(progQLinear); + if (progKVLinear) orion_release_program(progKVLinear); + free(hidden); + free(normed); + free(cpu_q); + free(cpu_k); + free(cpu_v); + free(ane_q); + free(ane_k); + free(ane_v); + free(cpu_attn); + free(ane_attn); + free(input_ln); + free(q_proj); + free(k_proj); + free(v_proj); + free(o_proj); + free(q_norm); + free(k_norm); + orion_qwen35_manifest_free(manifest); + return exit_code; + } +} diff --git a/tests/test_qwen35_9b_lora_ane_train_pairs.m b/tests/test_qwen35_9b_lora_ane_train_pairs.m new file mode 100644 index 0000000..48db37f --- /dev/null +++ b/tests/test_qwen35_9b_lora_ane_train_pairs.m @@ -0,0 +1,172 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +static NSArray *load_pairs(NSString *path) { + NSString *text = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:nil]; + if (!text) return nil; + NSMutableArray *items = [NSMutableArray array]; + [text enumerateLinesUsingBlock:^(NSString * _Nonnull line, BOOL * _Nonnull stop) { + (void)stop; + if (line.length == 0) return; + NSData *data = [line dataUsingEncoding:NSUTF8StringEncoding]; + NSDictionary *obj = [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; + if ([obj isKindOfClass:[NSDictionary class]]) { + [items addObject:obj]; + } + }]; + return items; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 5) { + fprintf(stderr, "usage: %s [resume_adapter_dir]\n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + NSString *pairsPath = [NSString stringWithUTF8String:argv[2]]; + const char *adapter_out_dir = argv[3]; + int max_steps = atoi(argv[4]); + const char *resume_adapter_dir = argc >= 6 ? argv[5] : NULL; + if (max_steps <= 0) { + fprintf(stderr, "FAIL: max_steps must be > 0\n"); + return 1; + } + + NSArray *pairs = load_pairs(pairsPath); + if (!pairs || pairs.count == 0) { + fprintf(stderr, "FAIL: pair file empty or unreadable\n"); + return 1; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + OrionQwen9BLoRATrainer trainer; + orion_qwen9b_lora_trainer_init(&trainer, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_qwen9b_lora_trainer_attach_ce_context(&trainer, embed_path, manifest)) { + fprintf(stderr, "FAIL: trainer CE context attach failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + int resume_loaded = 0; + int resume_start_step = trainer.step; + if (resume_adapter_dir && strlen(resume_adapter_dir) > 0) { + if (!orion_qwen9b_lora_trainer_load(&trainer, resume_adapter_dir)) { + fprintf(stderr, "FAIL: resume adapter load failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + resume_loaded = 1; + resume_start_step = trainer.step; + } + + double loss_sum = 0.0; + float loss_first = NAN; + float loss_last = NAN; + float loss_min = INFINITY; + float loss_max = -INFINITY; + double q_grad_last = 0.0; + double v_grad_last = 0.0; + double q_param_last = 0.0; + double v_param_last = 0.0; + int steps_completed = 0; + NSMutableSet *sampleIds = [NSMutableSet set]; + + for (NSDictionary *pair in pairs) { + if (steps_completed >= max_steps) break; + int inputToken = [pair[@"input_token"] intValue]; + int targetToken = [pair[@"target_token"] intValue]; + NSString *sampleId = pair[@"sample_id"] ?: @"unknown"; + [sampleIds addObject:sampleId]; + + OrionQwen9BLoRASmokeResult result; + if (!orion_qwen9b_lora_train_smoke1_ane_qv_base(blob_dir, &trainer, inputToken, targetToken, &result)) { + fprintf(stderr, "FAIL: ane pair step failed at step %d for sample %s\n", steps_completed + 1, sampleId.UTF8String); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (!isfinite(result.loss)) { + fprintf(stderr, "FAIL: non-finite ane loss at step %d\n", steps_completed + 1); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (steps_completed == 0) loss_first = result.loss; + loss_last = result.loss; + if (result.loss < loss_min) loss_min = result.loss; + if (result.loss > loss_max) loss_max = result.loss; + loss_sum += result.loss; + q_grad_last = result.q_grad_abs_sum; + v_grad_last = result.v_grad_abs_sum; + q_param_last = result.q_param_abs_sum; + v_param_last = result.v_param_abs_sum; + steps_completed += 1; + } + + if (steps_completed == 0) { + fprintf(stderr, "FAIL: zero ANE steps completed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (!orion_qwen9b_lora_trainer_save(&trainer, adapter_out_dir)) { + fprintf(stderr, "FAIL: ANE adapter save failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + OrionQwen9BLoRATrainer reloaded; + orion_qwen9b_lora_trainer_init(&reloaded, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + int reload_ok = orion_qwen9b_lora_trainer_load(&reloaded, adapter_out_dir) && + orion_qwen9b_lora_trainer_compare(&trainer, &reloaded, 1e-6f); + + printf("PASS: qwen35 9b lora ane pair training\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" pairs_path=%s\n", pairsPath.UTF8String); + printf(" adapter_out_dir=%s\n", adapter_out_dir); + printf(" layer_idx=%d\n", trainer.layer_idx); + printf(" rank=%d\n", trainer.q_proj.rank); + printf(" alpha=%.1f\n", trainer.q_proj.alpha); + printf(" steps_completed=%d\n", steps_completed); + printf(" samples_seen=%lu\n", (unsigned long)sampleIds.count); + printf(" resume_loaded=%d\n", resume_loaded); + printf(" resume_start_step=%d\n", resume_start_step); + printf(" trainer_step_final=%d\n", trainer.step); + printf(" loss_first=%.6f\n", loss_first); + printf(" loss_last=%.6f\n", loss_last); + printf(" loss_min=%.6f\n", loss_min); + printf(" loss_max=%.6f\n", loss_max); + printf(" loss_avg=%.6f\n", (float)(loss_sum / (double)steps_completed)); + printf(" q_grad_abs_sum_last=%.6f\n", q_grad_last); + printf(" v_grad_abs_sum_last=%.6f\n", v_grad_last); + printf(" q_param_abs_sum_last=%.6f\n", q_param_last); + printf(" v_param_abs_sum_last=%.6f\n", v_param_last); + printf(" reload_compare_ok=%d\n", reload_ok); + printf(" compile_cache_hit=%d\n", orion_qwen9b_lora_ane_train_bridge_last_compile_cache_hit()); + printf(" compile_cache_q_hit=%d\n", orion_qwen9b_lora_ane_train_bridge_last_q_cache_hit()); + printf(" compile_cache_kv_hit=%d\n", orion_qwen9b_lora_ane_train_bridge_last_kv_cache_hit()); + printf(" compile_cache_source=%s\n", orion_qwen9b_lora_ane_train_bridge_last_compile_cache_source()); + printf(" next_blocker=%s\n", "ane longer canary/preflight and cpu-vs-ane drift review"); + + orion_qwen9b_lora_trainer_free(&reloaded); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return reload_ok ? 0 : 1; + } +} diff --git a/tests/test_qwen35_9b_lora_ane_train_smoke1.m b/tests/test_qwen35_9b_lora_ane_train_smoke1.m new file mode 100644 index 0000000..bdc51fa --- /dev/null +++ b/tests/test_qwen35_9b_lora_ane_train_smoke1.m @@ -0,0 +1,407 @@ +#import +#import +#import +#include +#include +#include +#include +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" +#import "../kernels/training/stories_cpu_ops.h" +#import "../kernels/training/qwen_lora_train.h" + +static float abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return (float)total; +} + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static void expand_grouped_value(const float *v, + int n_head, + int n_kv_head, + int head_dim, + float *attn_cat) { + int q_per_kv = n_head / n_kv_head; + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + memcpy(attn_cat + (size_t)h * head_dim, + v + (size_t)kv_head * head_dim, + (size_t)head_dim * sizeof(float)); + } +} + +static void reduce_grouped_value_grad(const float *d_attn_cat, + int n_head, + int n_kv_head, + int head_dim, + float *d_v) { + int q_per_kv = n_head / n_kv_head; + memset(d_v, 0, (size_t)n_kv_head * head_dim * sizeof(float)); + for (int h = 0; h < n_head; h++) { + int kv_head = h / q_per_kv; + float *dst = d_v + (size_t)kv_head * head_dim; + const float *src = d_attn_cat + (size_t)h * head_dim; + for (int i = 0; i < head_dim; i++) dst[i] += src[i]; + } +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s [input_token] [target_token]\n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + const char *adapter_out_dir = argv[2]; + const int input_token = argc >= 4 ? atoi(argv[3]) : 27; + const int target_token = argc >= 5 ? atoi(argv[4]) : 91; + const int bucket = 32; + const int seq_len = 1; + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + if (!orion_ane_init()) { + fprintf(stderr, "FAIL: ane init failed\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + + const int layer_idx = manifest->n_layer - 1; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + NSString *blobDir = [NSString stringWithUTF8String:blob_dir]; + + OrionQwen9BLoRATrainer trainer; + orion_qwen9b_lora_trainer_init(&trainer, manifest, layer_idx, 8, 16.0f, 1e-3f, 1337u); + + float *hidden_in = (float *)calloc((size_t)d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *base_q = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *base_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *delta_q = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *delta_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *q_full = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *v_raw = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *attn_cat = (float *)calloc((size_t)q_dim, sizeof(float)); + float *gated = (float *)calloc((size_t)q_dim, sizeof(float)); + float *mixer = (float *)calloc((size_t)d_model, sizeof(float)); + float *hidden_mid = (float *)calloc((size_t)d_model, sizeof(float)); + float *post_norm = (float *)calloc((size_t)d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *hidden_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = (float *)calloc((size_t)d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_out = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_post_norm = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_mid = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_mixer = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_gated = (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_attn_cat = (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_gate_half = (float *)calloc((size_t)q_dim, sizeof(float)); + float *d_q_full = (float *)calloc((size_t)(q_dim * 2), sizeof(float)); + float *d_v = (float *)calloc((size_t)kv_dim, sizeof(float)); + float *d_normed_from_q = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_normed_from_v = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *d_hidden_mid_from_post = NULL; + float *d_post_ln_weight_grad = NULL; + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + char final_norm_path[2048], embed_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + float *final_norm_weight = orion_read_blob_f32_exact(final_norm_path, d_model); + IOSurfaceRef ioIn = NULL; + IOSurfaceRef ioQ = NULL; + IOSurfaceRef ioK = NULL; + IOSurfaceRef ioV = NULL; + OrionProgram *progQ = NULL; + OrionProgram *progKV = NULL; + NSString *milQ = nil; + NSString *milKV = nil; + int ok = 0; + float loss = NAN; + float q_grad_abs_sum = 0.0f; + float v_grad_abs_sum = 0.0f; + + if (!hidden_in || !normed || !base_q || !base_v || !delta_q || !delta_v || !q_full || !v_raw || + !attn_cat || !gated || !mixer || !hidden_mid || !post_norm || !mlp_out || !hidden_out || + !final_norm || !last_hidden || !d_last_hidden || !d_hidden_out || !d_post_norm || + !d_hidden_mid || !d_mixer || !d_gated || !d_attn_cat || !d_gate_half || !d_q_full || + !d_v || !d_normed_from_q || !d_normed_from_v || !d_normed || !input_ln || !post_ln || + !q_proj || !v_proj || !o_proj || !gate_proj || !up_proj || !down_proj || !final_norm_weight) { + fprintf(stderr, "FAIL: alloc or weight load failed\n"); + goto cleanup; + } + + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, input_token, hidden_in)) { + fprintf(stderr, "FAIL: frozen prefix hidden probe failed\n"); + goto cleanup; + } + orion_qwen_cpu_rmsnorm(hidden_in, input_ln, d_model, 1e-6f, normed); + + OrionModelConfig cfg = { + .n_layer = manifest->n_layer, + .n_head = manifest->n_head, + .n_kv_head = manifest->n_kv_head, + .d_model = manifest->d_model, + .head_dim = manifest->head_dim, + .hidden_dim = manifest->d_ff, + .vocab = manifest->vocab, + .max_seq = manifest->max_seq, + }; + milQ = compile_graph(orion_frontend_qwen35_prefill_q_proj(layer_idx, bucket, &cfg)); + milKV = compile_graph(orion_frontend_qwen35_prefill_kv_proj(layer_idx, bucket, &cfg)); + if (!milQ || !milKV) { + fprintf(stderr, "FAIL: mil graph build failed\n"); + goto cleanup; + } + progQ = orion_compile_mil(milQ.UTF8String, build_qproj_wdict(layer_idx, blobDir), "qwen35_9b_lora_ane_train_q"); + progKV = orion_compile_mil(milKV.UTF8String, build_kv_wdict(layer_idx, blobDir), "qwen35_9b_lora_ane_train_kv"); + if (!progQ || !progKV) { + fprintf(stderr, "FAIL: ane compile failed\n"); + goto cleanup; + } + + ioIn = make_cpu_seq_input_surface(hidden_in, seq_len, bucket, d_model); + ioQ = make_f32_surface((q_dim * 2) * bucket, 0.0f); + ioK = make_f32_surface(kv_dim * bucket, 0.0f); + ioV = make_f32_surface(kv_dim * bucket, 0.0f); + if (!ioIn || !ioQ || !ioK || !ioV) { + fprintf(stderr, "FAIL: iosurface alloc failed\n"); + goto cleanup; + } + + IOSurfaceRef ins[] = { ioIn }; + IOSurfaceRef outsQ[] = { ioQ }; + IOSurfaceRef outsKV[] = { ioK, ioV }; + if (!orion_eval(progQ, ins, 1, outsQ, 1) || !orion_eval(progKV, ins, 1, outsKV, 2)) { + fprintf(stderr, "FAIL: ane q/v eval failed\n"); + goto cleanup; + } + + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bucket, base_q); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bucket, base_v); + orion_qwen_lora_linear_delta_forward(normed, &trainer.q_proj, delta_q); + orion_qwen_lora_linear_delta_forward(normed, &trainer.v_proj, delta_v); + for (int i = 0; i < q_dim * 2; i++) q_full[i] = base_q[i] + delta_q[i]; + for (int i = 0; i < kv_dim; i++) v_raw[i] = base_v[i] + delta_v[i]; + + expand_grouped_value(v_raw, n_head, n_kv_head, head_dim, attn_cat); + for (int i = 0; i < q_dim; i++) { + float gate = 1.0f / (1.0f + expf(-q_full[q_dim + i])); + gated[i] = attn_cat[i] * gate; + } + cblas_sgemv(CblasRowMajor, CblasTrans, + d_model, q_dim, + 1.0f, o_proj, q_dim, gated, 1, + 0.0f, mixer, 1); + + for (int i = 0; i < d_model; i++) hidden_mid[i] = hidden_in[i] + mixer[i]; + orion_qwen_cpu_rmsnorm(hidden_mid, post_ln, d_model, 1e-6f, post_norm); + orion_qwen_cpu_swiglu_ffn(post_norm, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out); + for (int i = 0; i < d_model; i++) hidden_out[i] = hidden_mid[i] + mlp_out[i]; + orion_qwen_cpu_rmsnorm(hidden_out, final_norm_weight, d_model, 1e-6f, last_hidden); + + loss = orion_qwen_cpu_streaming_ce_tied_embedding(embed_path, last_hidden, d_model, manifest->vocab, target_token, d_last_hidden); + if (!isfinite(loss)) { + fprintf(stderr, "FAIL: non-finite loss\n"); + goto cleanup; + } + + orion_cpu_rmsnorm_bwd(d_hidden_out, final_norm, d_last_hidden, hidden_out, final_norm_weight, d_model, 1, 1e-6f); + orion_qwen_cpu_swiglu_ffn_bwd(post_norm, gate_proj, up_proj, down_proj, d_model, d_ff, d_hidden_out, d_post_norm); + d_hidden_mid_from_post = (float *)calloc((size_t)d_model, sizeof(float)); + d_post_ln_weight_grad = (float *)calloc((size_t)d_model, sizeof(float)); + orion_cpu_rmsnorm_bwd(d_hidden_mid_from_post, d_post_ln_weight_grad, d_post_norm, hidden_mid, post_ln, d_model, 1, 1e-6f); + for (int i = 0; i < d_model; i++) d_hidden_mid[i] = d_hidden_out[i] + d_hidden_mid_from_post[i]; + + memcpy(d_mixer, d_hidden_mid, (size_t)d_model * sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + d_model, q_dim, + 1.0f, o_proj, q_dim, d_mixer, 1, + 0.0f, d_gated, 1); + for (int i = 0; i < q_dim; i++) { + float gate_pre = q_full[q_dim + i]; + float gate = 1.0f / (1.0f + expf(-gate_pre)); + d_attn_cat[i] = d_gated[i] * gate; + d_gate_half[i] = d_gated[i] * attn_cat[i] * gate * (1.0f - gate); + d_q_full[q_dim + i] = d_gate_half[i]; + } + reduce_grouped_value_grad(d_attn_cat, n_head, n_kv_head, head_dim, d_v); + + orion_qwen_lora_linear_backward(normed, q_proj, &trainer.q_proj, d_q_full, d_normed_from_q); + orion_qwen_lora_linear_backward(normed, v_proj, &trainer.v_proj, d_v, d_normed_from_v); + for (int i = 0; i < d_model; i++) d_normed[i] = d_normed_from_q[i] + d_normed_from_v[i]; + + { + float *throwaway_weight_grad = (float *)calloc((size_t)d_model, sizeof(float)); + float *throwaway_dx = (float *)calloc((size_t)d_model, sizeof(float)); + orion_cpu_rmsnorm_bwd(throwaway_dx, throwaway_weight_grad, d_normed, hidden_in, input_ln, d_model, 1, 1e-6f); + free(throwaway_weight_grad); + free(throwaway_dx); + } + + q_grad_abs_sum = (float)orion_qwen_lora_grad_abs_sum(&trainer.q_proj); + v_grad_abs_sum = (float)orion_qwen_lora_grad_abs_sum(&trainer.v_proj); + orion_qwen9b_lora_trainer_step(&trainer); + if (!orion_qwen9b_lora_trainer_save(&trainer, adapter_out_dir)) { + fprintf(stderr, "FAIL: adapter save failed\n"); + goto cleanup; + } + + printf("PASS: qwen35 9b lora ane train smoke1\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" adapter_out_dir=%s\n", adapter_out_dir); + printf(" input_token=%d\n", input_token); + printf(" target_token=%d\n", target_token); + printf(" layer_idx=%d\n", layer_idx); + printf(" bucket=%d\n", bucket); + printf(" step=%d\n", trainer.step); + printf(" hidden_abs_sum=%.6f\n", abs_sum(hidden_in, d_model)); + printf(" normed_abs_sum=%.6f\n", abs_sum(normed, d_model)); + printf(" q_base_abs_sum=%.6f\n", abs_sum(base_q, q_dim * 2)); + printf(" v_base_abs_sum=%.6f\n", abs_sum(base_v, kv_dim)); + printf(" q_delta_abs_sum=%.6f\n", abs_sum(delta_q, q_dim * 2)); + printf(" v_delta_abs_sum=%.6f\n", abs_sum(delta_v, kv_dim)); + printf(" loss=%.6f\n", loss); + printf(" q_grad_abs_sum=%.6f\n", q_grad_abs_sum); + printf(" v_grad_abs_sum=%.6f\n", v_grad_abs_sum); + printf(" q_param_abs_sum=%.6f\n", + orion_qwen_lora_abs_sum(trainer.q_proj.a, trainer.q_proj.rank * trainer.q_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer.q_proj.b, trainer.q_proj.out_dim * trainer.q_proj.rank)); + printf(" v_param_abs_sum=%.6f\n", + orion_qwen_lora_abs_sum(trainer.v_proj.a, trainer.v_proj.rank * trainer.v_proj.in_dim) + + orion_qwen_lora_abs_sum(trainer.v_proj.b, trainer.v_proj.out_dim * trainer.v_proj.rank)); + printf(" next_blocker=%s\n", "adapter reload compare and smoke10 with ANE q/v base"); + ok = 1; + +cleanup: + if (ioIn) CFRelease(ioIn); + if (ioQ) CFRelease(ioQ); + if (ioK) CFRelease(ioK); + if (ioV) CFRelease(ioV); + if (progQ) orion_release_program(progQ); + if (progKV) orion_release_program(progKV); + free(hidden_in); free(normed); free(base_q); free(base_v); free(delta_q); free(delta_v); + free(q_full); free(v_raw); free(attn_cat); free(gated); free(mixer); free(hidden_mid); + free(post_norm); free(mlp_out); free(hidden_out); free(final_norm); free(last_hidden); + free(d_last_hidden); free(d_hidden_out); free(d_post_norm); free(d_hidden_mid); free(d_mixer); + free(d_gated); free(d_attn_cat); free(d_gate_half); free(d_q_full); free(d_v); free(d_normed_from_q); + free(d_normed_from_v); free(d_normed); free(d_hidden_mid_from_post); free(d_post_ln_weight_grad); + free(input_ln); free(post_ln); free(q_proj); free(v_proj); free(o_proj); free(gate_proj); free(up_proj); free(down_proj); + free(final_norm_weight); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return ok ? 0 : 1; + } +} diff --git a/tests/test_qwen35_9b_lora_reload_drift.m b/tests/test_qwen35_9b_lora_reload_drift.m new file mode 100644 index 0000000..03abf49 --- /dev/null +++ b/tests/test_qwen35_9b_lora_reload_drift.m @@ -0,0 +1,45 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + const char *blob_dir = argv[1]; + const char *adapter_dir = argv[2]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + OrionQwen9BLoRATrainer lhs; + OrionQwen9BLoRATrainer rhs; + orion_qwen9b_lora_trainer_init(&lhs, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + orion_qwen9b_lora_trainer_init(&rhs, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + int ok = orion_qwen9b_lora_trainer_load(&lhs, adapter_dir) && + orion_qwen9b_lora_trainer_load(&rhs, adapter_dir) && + orion_qwen9b_lora_trainer_compare(&lhs, &rhs, 1e-6f); + + printf("%s: qwen35 9b lora reload drift\n", ok ? "PASS" : "FAIL"); + printf(" blob_dir=%s\n", blob_dir); + printf(" adapter_dir=%s\n", adapter_dir); + printf(" step=%d\n", lhs.step); + printf(" layer_idx=%d\n", lhs.layer_idx); + printf(" q_param_abs_sum=%.6f\n", orion_qwen_lora_abs_sum(lhs.q_proj.a, lhs.q_proj.rank * lhs.q_proj.in_dim) + + orion_qwen_lora_abs_sum(lhs.q_proj.b, lhs.q_proj.out_dim * lhs.q_proj.rank)); + printf(" v_param_abs_sum=%.6f\n", orion_qwen_lora_abs_sum(lhs.v_proj.a, lhs.v_proj.rank * lhs.v_proj.in_dim) + + orion_qwen_lora_abs_sum(lhs.v_proj.b, lhs.v_proj.out_dim * lhs.v_proj.rank)); + printf(" compare_ok=%d\n", ok); + + orion_qwen9b_lora_trainer_free(&lhs); + orion_qwen9b_lora_trainer_free(&rhs); + orion_qwen35_manifest_free(manifest); + return ok ? 0 : 1; + } +} diff --git a/tests/test_qwen35_9b_lora_train_pairs.m b/tests/test_qwen35_9b_lora_train_pairs.m new file mode 100644 index 0000000..b0f1ea0 --- /dev/null +++ b/tests/test_qwen35_9b_lora_train_pairs.m @@ -0,0 +1,167 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +static NSArray *load_pairs(NSString *path) { + NSString *text = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:nil]; + if (!text) return nil; + NSMutableArray *items = [NSMutableArray array]; + [text enumerateLinesUsingBlock:^(NSString * _Nonnull line, BOOL * _Nonnull stop) { + if (line.length == 0) return; + NSData *data = [line dataUsingEncoding:NSUTF8StringEncoding]; + NSDictionary *obj = [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; + if ([obj isKindOfClass:[NSDictionary class]]) { + [items addObject:obj]; + } + }]; + return items; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 5) { + fprintf(stderr, "usage: %s [resume_adapter_dir]\n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + NSString *pairsPath = [NSString stringWithUTF8String:argv[2]]; + const char *adapter_out_dir = argv[3]; + int max_steps = atoi(argv[4]); + const char *resume_adapter_dir = argc >= 6 ? argv[5] : NULL; + if (max_steps <= 0) { + fprintf(stderr, "FAIL: max_steps must be > 0\n"); + return 1; + } + + NSArray *pairs = load_pairs(pairsPath); + if (!pairs || pairs.count == 0) { + fprintf(stderr, "FAIL: pair file empty or unreadable\n"); + return 1; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + OrionQwen9BLoRATrainer trainer; + orion_qwen9b_lora_trainer_init(&trainer, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_qwen9b_lora_trainer_attach_ce_context(&trainer, embed_path, manifest)) { + fprintf(stderr, "FAIL: trainer CE context attach failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + int resume_loaded = 0; + int resume_start_step = trainer.step; + if (resume_adapter_dir && strlen(resume_adapter_dir) > 0) { + if (!orion_qwen9b_lora_trainer_load(&trainer, resume_adapter_dir)) { + fprintf(stderr, "FAIL: resume adapter load failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + resume_loaded = 1; + resume_start_step = trainer.step; + } + + double loss_sum = 0.0; + float loss_first = NAN; + float loss_last = NAN; + float loss_min = INFINITY; + float loss_max = -INFINITY; + double q_grad_last = 0.0; + double v_grad_last = 0.0; + double q_param_last = 0.0; + double v_param_last = 0.0; + int steps_completed = 0; + NSMutableSet *sampleIds = [NSMutableSet set]; + + for (NSDictionary *pair in pairs) { + if (steps_completed >= max_steps) break; + int inputToken = [pair[@"input_token"] intValue]; + int targetToken = [pair[@"target_token"] intValue]; + NSString *sampleId = pair[@"sample_id"] ?: @"unknown"; + [sampleIds addObject:sampleId]; + + OrionQwen9BLoRASmokeResult result; + if (!orion_qwen9b_lora_train_smoke1(blob_dir, &trainer, inputToken, targetToken, &result)) { + fprintf(stderr, "FAIL: pair step failed at step %d for sample %s\n", steps_completed + 1, sampleId.UTF8String); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (!isfinite(result.loss)) { + fprintf(stderr, "FAIL: non-finite loss at step %d\n", steps_completed + 1); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (steps_completed == 0) loss_first = result.loss; + loss_last = result.loss; + if (result.loss < loss_min) loss_min = result.loss; + if (result.loss > loss_max) loss_max = result.loss; + loss_sum += result.loss; + q_grad_last = result.q_grad_abs_sum; + v_grad_last = result.v_grad_abs_sum; + q_param_last = result.q_param_abs_sum; + v_param_last = result.v_param_abs_sum; + steps_completed += 1; + } + + if (steps_completed == 0) { + fprintf(stderr, "FAIL: zero steps completed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (!orion_qwen9b_lora_trainer_save(&trainer, adapter_out_dir)) { + fprintf(stderr, "FAIL: adapter save failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + OrionQwen9BLoRATrainer reloaded; + orion_qwen9b_lora_trainer_init(&reloaded, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + int reload_ok = orion_qwen9b_lora_trainer_load(&reloaded, adapter_out_dir) && + orion_qwen9b_lora_trainer_compare(&trainer, &reloaded, 1e-6f); + + printf("PASS: qwen35 9b lora pair training\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" pairs_path=%s\n", pairsPath.UTF8String); + printf(" adapter_out_dir=%s\n", adapter_out_dir); + printf(" layer_idx=%d\n", trainer.layer_idx); + printf(" rank=%d\n", trainer.q_proj.rank); + printf(" alpha=%.1f\n", trainer.q_proj.alpha); + printf(" steps_completed=%d\n", steps_completed); + printf(" samples_seen=%lu\n", (unsigned long)sampleIds.count); + printf(" resume_loaded=%d\n", resume_loaded); + printf(" resume_start_step=%d\n", resume_start_step); + printf(" trainer_step_final=%d\n", trainer.step); + printf(" loss_first=%.6f\n", loss_first); + printf(" loss_last=%.6f\n", loss_last); + printf(" loss_min=%.6f\n", loss_min); + printf(" loss_max=%.6f\n", loss_max); + printf(" loss_avg=%.6f\n", (float)(loss_sum / (double)steps_completed)); + printf(" q_grad_abs_sum_last=%.6f\n", q_grad_last); + printf(" v_grad_abs_sum_last=%.6f\n", v_grad_last); + printf(" q_param_abs_sum_last=%.6f\n", q_param_last); + printf(" v_param_abs_sum_last=%.6f\n", v_param_last); + printf(" reload_compare_ok=%d\n", reload_ok); + printf(" next_blocker=%s\n", "crpg mainline launch gate and ANE hybrid parity residual review"); + + orion_qwen9b_lora_trainer_free(&reloaded); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return reload_ok ? 0 : 1; + } +} diff --git a/tests/test_qwen35_9b_lora_train_sequences.m b/tests/test_qwen35_9b_lora_train_sequences.m new file mode 100644 index 0000000..01615fa --- /dev/null +++ b/tests/test_qwen35_9b_lora_train_sequences.m @@ -0,0 +1,354 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +static NSArray *load_sequences(NSString *path) { + NSString *text = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:nil]; + if (!text) return nil; + NSMutableArray *items = [NSMutableArray array]; + [text enumerateLinesUsingBlock:^(NSString * _Nonnull line, BOOL * _Nonnull stop) { + (void)stop; + if (line.length == 0) return; + NSData *data = [line dataUsingEncoding:NSUTF8StringEncoding]; + NSDictionary *obj = [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; + if ([obj isKindOfClass:[NSDictionary class]]) { + [items addObject:obj]; + } + }]; + return items; +} + +static NSData *cached_hidden_for_token(NSMutableDictionary *cache, + const char *blob_dir, + const OrionQwen35Manifest *manifest, + int token_id) { + NSNumber *key = @(token_id); + NSData *hiddenData = cache[key]; + if (hiddenData) return hiddenData; + + size_t bytes = (size_t)manifest->d_model * sizeof(float); + float *hidden = (float *)calloc((size_t)manifest->d_model, sizeof(float)); + if (!hidden) return nil; + if (!orion_qwen9b_lora_frozen_prefix_hidden(blob_dir, manifest, token_id, hidden)) { + free(hidden); + return nil; + } + + hiddenData = [NSData dataWithBytesNoCopy:hidden length:bytes freeWhenDone:YES]; + if (!hiddenData) { + free(hidden); + return nil; + } + cache[key] = hiddenData; + return hiddenData; +} + +static NSMutableDictionary *load_hidden_cache_dir(NSString *cacheDir, + const OrionQwen35Manifest *manifest) { + if (!cacheDir) return nil; + NSString *indexPath = [cacheDir stringByAppendingPathComponent:@"hidden_cache_index.json"]; + NSData *indexData = [NSData dataWithContentsOfFile:indexPath]; + if (!indexData) return nil; + NSDictionary *index = [NSJSONSerialization JSONObjectWithData:indexData options:0 error:nil]; + if (![index isKindOfClass:[NSDictionary class]]) return nil; + if ([index[@"d_model"] intValue] != manifest->d_model) return nil; + + NSString *binPath = index[@"bin_path"]; + NSData *binData = [NSData dataWithContentsOfFile:binPath options:NSDataReadingMappedIfSafe error:nil]; + NSArray *tokens = index[@"tokens"]; + NSArray *offsets = index[@"offsets"]; + NSArray *counts = index[@"counts"]; + if (!binData || ![tokens isKindOfClass:[NSArray class]] || + ![offsets isKindOfClass:[NSArray class]] || ![counts isKindOfClass:[NSArray class]] || + tokens.count != offsets.count || tokens.count != counts.count) { + return nil; + } + + NSMutableDictionary *cache = [NSMutableDictionary dictionaryWithCapacity:tokens.count]; + for (NSUInteger idx = 0; idx < tokens.count; idx++) { + NSNumber *token = tokens[idx]; + NSUInteger offset = (NSUInteger)[offsets[idx] unsignedLongLongValue]; + NSUInteger count = (NSUInteger)[counts[idx] unsignedLongLongValue]; + NSUInteger bytes = count * sizeof(float); + if (offset + bytes > binData.length) return nil; + NSData *hiddenData = [binData subdataWithRange:NSMakeRange(offset, bytes)]; + cache[token] = hiddenData; + } + return cache; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 5) { + fprintf(stderr, "usage: %s [hidden_cache_dir] [grad_accum_sequences] [sequence_bucket] [resume_adapter_dir]\n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + NSString *sequencesPath = [NSString stringWithUTF8String:argv[2]]; + const char *adapter_out_dir = argv[3]; + int max_sequences = atoi(argv[4]); + NSString *hiddenCacheDir = argc >= 6 ? [NSString stringWithUTF8String:argv[5]] : nil; + int grad_accum_sequences = argc >= 7 ? atoi(argv[6]) : 4; + int sequence_bucket = argc >= 8 ? atoi(argv[7]) : 64; + const char *resume_adapter_dir = argc >= 9 ? argv[8] : NULL; + if (max_sequences <= 0) { + fprintf(stderr, "FAIL: max_sequences must be > 0\n"); + return 1; + } + if (grad_accum_sequences <= 0) { + fprintf(stderr, "FAIL: grad_accum_sequences must be > 0\n"); + return 1; + } + if (sequence_bucket <= 0) { + fprintf(stderr, "FAIL: sequence_bucket must be > 0\n"); + return 1; + } + + NSArray *sequences = load_sequences(sequencesPath); + if (!sequences || sequences.count == 0) { + fprintf(stderr, "FAIL: sequence file empty or unreadable\n"); + return 1; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + OrionQwen9BLoRATrainer trainer; + orion_qwen9b_lora_trainer_init(&trainer, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_qwen9b_lora_trainer_attach_ce_context(&trainer, embed_path, manifest)) { + fprintf(stderr, "FAIL: trainer CE context attach failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (!orion_qwen9b_lora_trainer_attach_cpu_train_context(&trainer, blob_dir, manifest)) { + fprintf(stderr, "FAIL: trainer CPU train context attach failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + int resume_loaded = 0; + int resume_start_step = trainer.step; + if (resume_adapter_dir && strlen(resume_adapter_dir) > 0) { + if (!orion_qwen9b_lora_trainer_load(&trainer, resume_adapter_dir)) { + fprintf(stderr, "FAIL: sequence resume adapter load failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + resume_loaded = 1; + resume_start_step = trainer.step; + } + + double loss_sum = 0.0; + float loss_first = NAN; + float loss_last = NAN; + float loss_min = INFINITY; + float loss_max = -INFINITY; + double q_grad_last = 0.0; + double v_grad_last = 0.0; + double q_param_last = 0.0; + double v_param_last = 0.0; + int sequences_completed = 0; + int pairs_seen = 0; + int tokens_seen = 0; + int packed_batches_completed = 0; + int packed_batch_tokens_last = 0; + int packed_batch_tokens_max = 0; + int packed_sequences_last = 0; + int packed_sequences_max = 0; + NSMutableSet *sampleIds = [NSMutableSet set]; + NSMutableDictionary *hiddenCache = load_hidden_cache_dir(hiddenCacheDir, manifest); + if (!hiddenCache) hiddenCache = [NSMutableDictionary dictionary]; + NSUInteger hiddenCachePrefilled = hiddenCache.count; + int hidden_cache_hits = 0; + int hidden_cache_misses = 0; + + NSUInteger sequenceIndex = 0; + while (sequenceIndex < sequences.count && sequences_completed < max_sequences) { + NSMutableArray *batchSequences = [NSMutableArray arrayWithCapacity:(NSUInteger)grad_accum_sequences]; + int batchPairs = 0; + int batchTokens = 0; + + while (sequenceIndex < sequences.count && + sequences_completed + (int)batchSequences.count < max_sequences && + (int)batchSequences.count < grad_accum_sequences) { + NSDictionary *sequence = sequences[sequenceIndex]; + sequenceIndex += 1; + NSArray *tokenIds = sequence[@"token_ids"]; + if (![tokenIds isKindOfClass:[NSArray class]] || tokenIds.count < 2) continue; + int seq_pairs = (int)tokenIds.count - 1; + if (batchPairs > 0 && batchPairs + seq_pairs > sequence_bucket) { + sequenceIndex -= 1; + break; + } + [batchSequences addObject:sequence]; + batchPairs += seq_pairs; + batchTokens += (int)tokenIds.count; + if (batchPairs >= sequence_bucket) break; + } + + if (batchSequences.count == 0 || batchPairs <= 0) { + continue; + } + + const float **hiddenBatch = (const float **)calloc((size_t)batchPairs, sizeof(float *)); + int *targetBatch = (int *)calloc((size_t)batchPairs, sizeof(int)); + if (!hiddenBatch || !targetBatch) { + free(hiddenBatch); + free(targetBatch); + fprintf(stderr, "FAIL: sequence batch alloc failed at sequence %d\n", sequences_completed + 1); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + int writeIdx = 0; + NSString *sampleLabel = @"packed-batch"; + for (NSDictionary *sequence in batchSequences) { + NSArray *tokenIds = sequence[@"token_ids"]; + NSString *sampleId = sequence[@"sample_id"] ?: @"unknown"; + if (sampleLabel.length == 0 || [sampleLabel isEqualToString:@"packed-batch"]) { + sampleLabel = sampleId; + } + [sampleIds addObject:sampleId]; + for (NSUInteger idx = 0; idx + 1 < tokenIds.count; idx++) { + int inputToken = [tokenIds[idx] intValue]; + int targetToken = [tokenIds[idx + 1] intValue]; + int hadHidden = hiddenCache[@(inputToken)] != nil; + NSData *hiddenData = cached_hidden_for_token(hiddenCache, blob_dir, manifest, inputToken); + if (!hiddenData) { + free(hiddenBatch); + free(targetBatch); + fprintf(stderr, "FAIL: hidden cache build failed at sequence %d for token %d\n", + sequences_completed + 1, inputToken); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (hadHidden) hidden_cache_hits += 1; + else hidden_cache_misses += 1; + hiddenBatch[writeIdx] = (const float *)hiddenData.bytes; + targetBatch[writeIdx] = targetToken; + writeIdx += 1; + } + } + + OrionQwen9BLoRABatchResult batchResult; + if (!orion_qwen9b_lora_train_hidden_batch(blob_dir, manifest, &trainer, + hiddenBatch, targetBatch, batchPairs, &batchResult)) { + free(hiddenBatch); + free(targetBatch); + fprintf(stderr, "FAIL: sequence packed batch failed at sequence %d for sample %s\n", + sequences_completed + 1, sampleLabel.UTF8String); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + free(hiddenBatch); + free(targetBatch); + + if (batchResult.items_completed != batchPairs || !isfinite(batchResult.loss_avg)) { + fprintf(stderr, "FAIL: invalid sequence packed batch result at sequence %d\n", sequences_completed + 1); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + float seq_loss_avg = batchResult.loss_avg; + if (sequences_completed == 0) loss_first = seq_loss_avg; + loss_last = seq_loss_avg; + if (seq_loss_avg < loss_min) loss_min = seq_loss_avg; + if (seq_loss_avg > loss_max) loss_max = seq_loss_avg; + loss_sum += batchResult.loss_avg * (double)batchSequences.count; + q_grad_last = batchResult.q_grad_abs_sum_last; + v_grad_last = batchResult.v_grad_abs_sum_last; + q_param_last = batchResult.q_param_abs_sum_last; + v_param_last = batchResult.v_param_abs_sum_last; + pairs_seen += batchPairs; + packed_batches_completed += 1; + packed_batch_tokens_last = batchPairs; + if (batchPairs > packed_batch_tokens_max) packed_batch_tokens_max = batchPairs; + packed_sequences_last = (int)batchSequences.count; + if ((int)batchSequences.count > packed_sequences_max) packed_sequences_max = (int)batchSequences.count; + + sequences_completed += (int)batchSequences.count; + tokens_seen += batchTokens; + fprintf(stderr, "INFO: completed_sequence=%d/%d packed_sequences=%lu sample_id=%s tokens=%d packed_pairs=%d avg_loss=%.6f\n", + sequences_completed, max_sequences, (unsigned long)batchSequences.count, sampleLabel.UTF8String, + batchTokens, batchPairs, seq_loss_avg); + fflush(stderr); + } + + if (sequences_completed == 0) { + fprintf(stderr, "FAIL: zero sequences completed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (!orion_qwen9b_lora_trainer_save(&trainer, adapter_out_dir)) { + fprintf(stderr, "FAIL: sequence adapter save failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + OrionQwen9BLoRATrainer reloaded; + orion_qwen9b_lora_trainer_init(&reloaded, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + int reload_ok = orion_qwen9b_lora_trainer_load(&reloaded, adapter_out_dir) && + orion_qwen9b_lora_trainer_compare(&trainer, &reloaded, 1e-6f); + + printf("PASS: qwen35 9b lora sequence training\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" sequences_path=%s\n", sequencesPath.UTF8String); + printf(" adapter_out_dir=%s\n", adapter_out_dir); + printf(" hidden_cache_dir=%s\n", hiddenCacheDir ? hiddenCacheDir.UTF8String : ""); + printf(" layer_idx=%d\n", trainer.layer_idx); + printf(" rank=%d\n", trainer.q_proj.rank); + printf(" alpha=%.1f\n", trainer.q_proj.alpha); + printf(" grad_accum_sequences=%d\n", grad_accum_sequences); + printf(" sequence_bucket=%d\n", sequence_bucket); + printf(" sequences_completed=%d\n", sequences_completed); + printf(" samples_seen=%lu\n", (unsigned long)sampleIds.count); + printf(" resume_loaded=%d\n", resume_loaded); + printf(" resume_start_step=%d\n", resume_start_step); + printf(" trainer_step_final=%d\n", trainer.step); + printf(" tokens_seen=%d\n", tokens_seen); + printf(" pairs_seen=%d\n", pairs_seen); + printf(" hidden_cache_prefilled=%lu\n", (unsigned long)hiddenCachePrefilled); + printf(" cached_hidden_tokens=%lu\n", (unsigned long)hiddenCache.count); + printf(" hidden_cache_hits=%d\n", hidden_cache_hits); + printf(" hidden_cache_misses=%d\n", hidden_cache_misses); + printf(" packed_batches_completed=%d\n", packed_batches_completed); + printf(" packed_sequences_last=%d\n", packed_sequences_last); + printf(" packed_sequences_max=%d\n", packed_sequences_max); + printf(" packed_batch_tokens_last=%d\n", packed_batch_tokens_last); + printf(" packed_batch_tokens_max=%d\n", packed_batch_tokens_max); + printf(" avg_tokens_per_sequence=%.3f\n", (double)tokens_seen / (double)sequences_completed); + printf(" loss_first=%.6f\n", loss_first); + printf(" loss_last=%.6f\n", loss_last); + printf(" loss_min=%.6f\n", loss_min); + printf(" loss_max=%.6f\n", loss_max); + printf(" loss_avg=%.6f\n", (float)(loss_sum / (double)sequences_completed)); + printf(" q_grad_abs_sum_last=%.6f\n", q_grad_last); + printf(" v_grad_abs_sum_last=%.6f\n", v_grad_last); + printf(" q_param_abs_sum_last=%.6f\n", q_param_last); + printf(" v_param_abs_sum_last=%.6f\n", v_param_last); + printf(" reload_compare_ok=%d\n", reload_ok); + printf(" next_blocker=%s\n", "validation split and hybrid parity gate promotion"); + + orion_qwen9b_lora_trainer_free(&reloaded); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return reload_ok ? 0 : 1; + } +} diff --git a/tests/test_qwen35_9b_lora_train_smoke1.m b/tests/test_qwen35_9b_lora_train_smoke1.m new file mode 100644 index 0000000..5f055d7 --- /dev/null +++ b/tests/test_qwen35_9b_lora_train_smoke1.m @@ -0,0 +1,79 @@ +#import +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/training/qwen_lora_train.h" + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + const char *adapter_out_dir = argv[2]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest load failed\n"); + return 1; + } + + OrionQwen9BLoRATrainer trainer; + orion_qwen9b_lora_trainer_init(&trainer, manifest, manifest->n_layer - 1, 8, 16.0f, 1e-3f, 1337u); + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_qwen9b_lora_trainer_attach_ce_context(&trainer, embed_path, manifest)) { + fprintf(stderr, "FAIL: trainer CE context attach failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + OrionQwen9BLoRASmokeResult result; + if (!orion_qwen9b_lora_train_smoke1(blob_dir, &trainer, 0, 1, &result)) { + fprintf(stderr, "FAIL: smoke1 train step failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + if (!isfinite(result.loss) || result.loss <= 0.0f) { + fprintf(stderr, "FAIL: non-finite loss %.6f\n", result.loss); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (result.q_param_abs_sum <= 0.0 || result.v_param_abs_sum <= 0.0) { + fprintf(stderr, "FAIL: non-positive parameter abs sums\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (!orion_qwen9b_lora_trainer_save(&trainer, adapter_out_dir)) { + fprintf(stderr, "FAIL: adapter save failed\n"); + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 1; + } + + printf("PASS: qwen35 9b lora train smoke1\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" adapter_out_dir=%s\n", adapter_out_dir); + printf(" layer_idx=%d\n", trainer.layer_idx); + printf(" rank=%d\n", trainer.q_proj.rank); + printf(" alpha=%.1f\n", trainer.q_proj.alpha); + printf(" step=%d\n", trainer.step); + printf(" loss=%.6f\n", result.loss); + printf(" q_grad_abs_sum=%.6f\n", result.q_grad_abs_sum); + printf(" v_grad_abs_sum=%.6f\n", result.v_grad_abs_sum); + printf(" q_param_abs_sum=%.6f\n", result.q_param_abs_sum); + printf(" v_param_abs_sum=%.6f\n", result.v_param_abs_sum); + printf(" next_blocker=%s\n", "smoke10 and adapter reload drift"); + + orion_qwen9b_lora_trainer_free(&trainer); + orion_qwen35_manifest_free(manifest); + return 0; + } +} diff --git a/tests/test_qwen35_9b_manifest_loader.m b/tests/test_qwen35_9b_manifest_loader.m new file mode 100644 index 0000000..94d01ac --- /dev/null +++ b/tests/test_qwen35_9b_manifest_loader.m @@ -0,0 +1,63 @@ +#import +#import +#import "../model/weight_loader.h" + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(argv[1]); + if (!manifest) { + fprintf(stderr, "FAIL: qwen35 9b manifest loader returned NULL\n"); + return 1; + } + + printf("PASS: qwen35 9b manifest loader\n"); + printf(" blob_dir=%s\n", manifest->blob_dir); + printf(" n_layer=%d\n", manifest->n_layer); + printf(" d_model=%d\n", manifest->d_model); + printf(" d_ff=%d\n", manifest->d_ff); + printf(" n_head=%d\n", manifest->n_head); + printf(" n_kv_head=%d\n", manifest->n_kv_head); + printf(" head_dim=%d\n", manifest->head_dim); + printf(" vocab=%d\n", manifest->vocab); + printf(" max_seq=%d\n", manifest->max_seq); + printf(" tie_word_embeddings=%d\n", manifest->tie_word_embeddings); + printf(" n_entries=%d\n", manifest->n_entries); + printf(" n_alias_entries=%d\n", manifest->n_alias_entries); + printf(" n_linear_layers=%d\n", manifest->n_linear_layers); + printf(" n_full_layers=%d\n", manifest->n_full_layers); + + if (manifest->n_layer != 32 || manifest->d_model != 4096 || manifest->d_ff != 12288) { + fprintf(stderr, "FAIL: unexpected 9b core dimensions\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (manifest->n_head != 16 || manifest->n_kv_head != 4 || manifest->head_dim != 256) { + fprintf(stderr, "FAIL: unexpected 9b head geometry\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (manifest->vocab != 248320 || manifest->max_seq != 262144) { + fprintf(stderr, "FAIL: unexpected 9b vocab/max_seq\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (manifest->n_linear_layers != 24 || manifest->n_full_layers != 8) { + fprintf(stderr, "FAIL: unexpected 9b hybrid topology\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (manifest->n_entries < 400) { + fprintf(stderr, "FAIL: too few exported entries for 9b\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + + orion_qwen35_manifest_free(manifest); + return 0; + } +} diff --git a/tests/test_qwen35_9b_prefill_cpu_decode_bridge.m b/tests/test_qwen35_9b_prefill_cpu_decode_bridge.m new file mode 100644 index 0000000..cb2c924 --- /dev/null +++ b/tests/test_qwen35_9b_prefill_cpu_decode_bridge.m @@ -0,0 +1,304 @@ +#import +#import +#include +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" +#include "model/weight_loader.h" +#include "kernels/inference/qwen_cpu_ops.h" + +static const OrionModelConfig kQwen35_9B = { + .n_layer = 32, + .n_head = 16, + .n_kv_head = 4, + .d_model = 4096, + .head_dim = 256, + .hidden_dim = 12288, + .vocab = 248320, + .max_seq = 262144, +}; + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_cpu_seq_input_surface(const float *x_seq, int seq_len, int bucket, int d_model) { + IOSurfaceRef s = make_f32_surface(d_model * bucket, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + ptr[c * bucket + t] = x_seq[t * d_model + c]; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static void read_ane_surface_prefix(IOSurfaceRef s, int channels, int seq_len, int bucket, float *out_seq) { + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < channels; c++) { + out_seq[t * channels + c] = ptr[c * bucket + t]; + } + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) dict[mil_path] = @{@"offset": @0, @"data": data}; +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_down_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_down_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static float *load_layer_exact(const char *blob_dir, int layer_idx, const char *name, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, name); + return orion_read_blob_f32_exact(path, count); +} + +static double mean_abs_diff(const float *a, const float *b, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)a[i] - (double)b[i]); + return total / (double)n; +} + +static double max_abs_diff(const float *a, const float *b, int n) { + double best = 0.0; + for (int i = 0; i < n; i++) { + double d = fabs((double)a[i] - (double)b[i]); + if (d > best) best = d; + } + return best; +} + +static double abs_sum_vec(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +int main(int argc, char **argv) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + if (!orion_ane_init()) { + fprintf(stderr, "FAIL: orion_ane_init failed\n"); + return 3; + } + + const char *blob_dir = argv[1]; + NSString *blob_dir_ns = @(blob_dir); + const int layer = 3; + const int seq_len = 2; + const int bucket = 32; + const int d_model = kQwen35_9B.d_model; + const int d_ff = kQwen35_9B.hidden_dim; + const int n_head = kQwen35_9B.n_head; + const int n_kv_head = kQwen35_9B.n_kv_head; + const int head_dim = kQwen35_9B.head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + NSString *mil_q = compile_graph(orion_frontend_qwen35_prefill_q_proj(layer, bucket, &kQwen35_9B)); + NSString *mil_kv = compile_graph(orion_frontend_qwen35_prefill_kv_proj(layer, bucket, &kQwen35_9B)); + NSString *mil_ffn = compile_graph(orion_frontend_qwen35_prefill_ffn(layer, bucket, &kQwen35_9B)); + if (!mil_q || !mil_kv || !mil_ffn) { + fprintf(stderr, "FAIL: compile_graph returned nil\n"); + return 4; + } + + OrionProgram *prog_q = orion_compile_mil(mil_q.UTF8String, build_qproj_wdict(layer, blob_dir_ns), "qwen35_9b_bridge_q"); + OrionProgram *prog_kv = orion_compile_mil(mil_kv.UTF8String, build_kv_wdict(layer, blob_dir_ns), "qwen35_9b_bridge_kv"); + OrionProgram *prog_ffn = orion_compile_mil(mil_ffn.UTF8String, build_ffn_wdict(layer, blob_dir_ns), "qwen35_9b_bridge_ffn"); + if (!prog_q || !prog_kv || !prog_ffn) { + fprintf(stderr, "FAIL: ANE program compile failed\n"); + if (prog_q) orion_release_program(prog_q); + if (prog_kv) orion_release_program(prog_kv); + if (prog_ffn) orion_release_program(prog_ffn); + return 5; + } + + float *x_seq = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + for (int t = 0; t < seq_len; t++) { + for (int c = 0; c < d_model; c++) { + x_seq[t * d_model + c] = sinf((float)(c * 0.001 + t * 0.01)) * 0.1f; + } + } + + float *input_ln = load_layer_exact(blob_dir, layer, "input_layernorm.bin", d_model); + float *post_ln = load_layer_exact(blob_dir, layer, "post_attention_layernorm.bin", d_model); + float *q_proj = load_layer_exact(blob_dir, layer, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_layer_exact(blob_dir, layer, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_layer_exact(blob_dir, layer, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_layer_exact(blob_dir, layer, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_layer_exact(blob_dir, layer, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_layer_exact(blob_dir, layer, "self_attn_k_norm.bin", head_dim); + float *gate_proj = load_layer_exact(blob_dir, layer, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_layer_exact(blob_dir, layer, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_layer_exact(blob_dir, layer, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm || + !gate_proj || !up_proj || !down_proj) { + fprintf(stderr, "FAIL: missing layer weights\n"); + return 6; + } + + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *cpu_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *cpu_hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *cpu_normed_post = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *cpu_mlp = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *cpu_hidden_final = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + + for (int t = 0; t < seq_len; t++) { + orion_qwen_cpu_rmsnorm(x_seq + t * d_model, input_ln, d_model, 1e-6f, normed + t * d_model); + } + orion_qwen_cpu_full_attention_prefill_with_rope(normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, 10000000.0f, 0.25f, cpu_attn); + for (int i = 0; i < seq_len * d_model; i++) cpu_hidden_attn[i] = x_seq[i] + cpu_attn[i]; + for (int t = 0; t < seq_len; t++) { + orion_qwen_cpu_rmsnorm(cpu_hidden_attn + t * d_model, post_ln, d_model, 1e-6f, cpu_normed_post + t * d_model); + orion_qwen_cpu_swiglu_ffn(cpu_normed_post + t * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, + cpu_mlp + t * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) cpu_hidden_final[i] = cpu_hidden_attn[i] + cpu_mlp[i]; + + IOSurfaceRef ioIn = make_cpu_seq_input_surface(x_seq, seq_len, bucket, d_model); + IOSurfaceRef ioQ = make_f32_surface((q_dim * 2) * bucket, 0.0f); + IOSurfaceRef ioK = make_f32_surface(kv_dim * bucket, 0.0f); + IOSurfaceRef ioV = make_f32_surface(kv_dim * bucket, 0.0f); + IOSurfaceRef ioHidden = make_f32_surface(d_model * bucket, 0.0f); + + IOSurfaceRef insQ[] = {ioIn}; + IOSurfaceRef outsQ[] = {ioQ}; + IOSurfaceRef outsKV[] = {ioK, ioV}; + if (!orion_eval(prog_q, insQ, 1, outsQ, 1) || !orion_eval(prog_kv, insQ, 1, outsKV, 2)) { + fprintf(stderr, "FAIL: ANE q/kv eval failed\n"); + return 7; + } + + float *q_proj_seq = (float *)calloc((size_t)seq_len * (q_dim * 2), sizeof(float)); + float *k_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + float *v_proj_seq = (float *)calloc((size_t)seq_len * kv_dim, sizeof(float)); + float *bridge_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *bridge_hidden_attn = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *bridge_hidden_final = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + + read_ane_surface_prefix(ioQ, q_dim * 2, seq_len, bucket, q_proj_seq); + read_ane_surface_prefix(ioK, kv_dim, seq_len, bucket, k_proj_seq); + read_ane_surface_prefix(ioV, kv_dim, seq_len, bucket, v_proj_seq); + + orion_qwen_cpu_full_attention_from_projections_with_rope(q_proj_seq, k_proj_seq, v_proj_seq, seq_len, + o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, + 10000000.0f, 0.25f, bridge_attn); + for (int i = 0; i < seq_len * d_model; i++) bridge_hidden_attn[i] = x_seq[i] + bridge_attn[i]; + + IOSurfaceRef ioFfnIn = make_cpu_seq_input_surface(bridge_hidden_attn, seq_len, bucket, d_model); + IOSurfaceRef insFFN[] = {ioFfnIn}; + IOSurfaceRef outsFFN[] = {ioHidden}; + if (!orion_eval(prog_ffn, insFFN, 1, outsFFN, 1)) { + fprintf(stderr, "FAIL: ANE ffn eval failed\n"); + return 8; + } + read_ane_surface_prefix(ioHidden, d_model, seq_len, bucket, bridge_hidden_final); + + double mean_diff = mean_abs_diff(cpu_hidden_final, bridge_hidden_final, seq_len * d_model); + double max_diff = max_abs_diff(cpu_hidden_final, bridge_hidden_final, seq_len * d_model); + printf("qwen35_9b_prefill_cpu_decode_bridge:\n"); + printf(" seq_len=%d\n", seq_len); + printf(" layer=%d\n", layer); + printf(" mean_abs_diff=%.6f\n", mean_diff); + printf(" max_abs_diff=%.6f\n", max_diff); + printf(" bridge_hidden_abs_sum=%.6f\n", abs_sum_vec(bridge_hidden_final, seq_len * d_model)); + + bool pass = isfinite(mean_diff) && isfinite(max_diff) && mean_diff < 0.05 && max_diff < 0.5; + printf(" status=%s\n", pass ? "PASS" : "FAIL"); + + CFRelease(ioIn); + CFRelease(ioQ); + CFRelease(ioK); + CFRelease(ioV); + CFRelease(ioHidden); + CFRelease(ioFfnIn); + orion_release_program(prog_q); + orion_release_program(prog_kv); + orion_release_program(prog_ffn); + + free(x_seq); + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(gate_proj); free(up_proj); free(down_proj); + free(normed); free(cpu_attn); free(cpu_hidden_attn); free(cpu_normed_post); free(cpu_mlp); free(cpu_hidden_final); + free(q_proj_seq); free(k_proj_seq); free(v_proj_seq); free(bridge_attn); free(bridge_hidden_attn); free(bridge_hidden_final); + + return pass ? 0 : 1; + } +} diff --git a/tests/test_qwen35_9b_prefill_frontend.m b/tests/test_qwen35_9b_prefill_frontend.m new file mode 100644 index 0000000..6d41a94 --- /dev/null +++ b/tests/test_qwen35_9b_prefill_frontend.m @@ -0,0 +1,90 @@ +// tests/test_qwen35_9b_prefill_frontend.m — Qwen3.5-9B ANE frontend smoke + +#import +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pass_uniform_outputs.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" + +static int g_passed = 0, g_failed = 0; + +#define TEST(name) do { \ + printf(" %-40s ", #name); \ + if (test_##name()) { printf("PASS\n"); g_passed++; } \ + else { printf("FAIL\n"); g_failed++; } \ +} while(0) + +static const OrionModelConfig kQwen35_9B = { + .n_layer = 32, + .n_head = 16, + .n_kv_head = 4, + .d_model = 4096, + .head_dim = 256, + .hidden_dim = 12288, + .vocab = 248320, + .max_seq = 262144, +}; + +static NSString* compile_graph(OrionGraph* g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString* mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static bool test_q_proj(void) { + OrionGraph* g = orion_frontend_qwen35_prefill_q_proj(3, 64, &kQwen35_9B); + NSString* mil = compile_graph(g); + if (!mil) return false; + return [mil containsString:@"self_attn_q_proj.bin"] && + [mil containsString:@"input_layernorm.bin"] && + [mil containsString:@"q_proj"]; +} + +static bool test_kv_proj(void) { + OrionGraph* g = orion_frontend_qwen35_prefill_kv_proj(3, 64, &kQwen35_9B); + bool needs_uniform_pad = orion_pass_uniform_outputs(g); + if (needs_uniform_pad) { + NSLog(@"kv_proj unexpectedly needs uniform output padding"); + orion_graph_free(g); + return false; + } + NSString* mil = compile_graph(g); + if (!mil) return false; + return [mil containsString:@"self_attn_k_proj.bin"] && + [mil containsString:@"self_attn_v_proj.bin"] && + [mil containsString:@"k_proj"] && + [mil containsString:@"v_proj"]; +} + +static bool test_ffn(void) { + OrionGraph* g = orion_frontend_qwen35_prefill_ffn(3, 64, &kQwen35_9B); + NSString* mil = compile_graph(g); + if (!mil) return false; + return [mil containsString:@"post_attention_layernorm.bin"] && + [mil containsString:@"mlp_gate_proj.bin"] && + [mil containsString:@"mlp_up_proj.bin"] && + [mil containsString:@"mlp_down_proj.bin"] && + [mil containsString:@"tanh("] && + [mil containsString:@"_sig"]; +} + +int main(int argc __attribute__((unused)), char* argv[] __attribute__((unused))) { + @autoreleasepool { + printf("test_qwen35_9b_prefill_frontend:\n"); + TEST(q_proj); + TEST(kv_proj); + TEST(ffn); + printf("\n%d/%d passed\n", g_passed, g_passed + g_failed); + return g_failed > 0 ? 1 : 0; + } +} diff --git a/tests/test_qwen35_9b_prefill_runtime_smoke.m b/tests/test_qwen35_9b_prefill_runtime_smoke.m new file mode 100644 index 0000000..46ad3ef --- /dev/null +++ b/tests/test_qwen35_9b_prefill_runtime_smoke.m @@ -0,0 +1,232 @@ +// tests/test_qwen35_9b_prefill_runtime_smoke.m +// Qwen3.5-9B ANE prefill runtime smoke: compile + eval q_proj / kv_proj / ffn. + +#import +#import +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" +#include "core/ane_runtime.h" + +static int g_passed = 0, g_failed = 0; + +#define TEST(name) do { \ + printf(" %-40s ", #name); \ + if (test_##name(blob_dir)) { printf("PASS\n"); g_passed++; } \ + else { printf("FAIL\n"); g_failed++; } \ +} while(0) + +static const OrionModelConfig kQwen35_9B = { + .n_layer = 32, + .n_head = 16, + .n_kv_head = 4, + .d_model = 4096, + .head_dim = 256, + .hidden_dim = 12288, + .vocab = 248320, + .max_seq = 262144, +}; + +static IOSurfaceRef make_f32_surface(int count, float fill) { + size_t bytes = (size_t)count * sizeof(float); + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef)@{ + (id)kIOSurfaceWidth: @(bytes), + (id)kIOSurfaceHeight: @1, + (id)kIOSurfaceBytesPerElement: @1, + (id)kIOSurfaceBytesPerRow: @(bytes), + (id)kIOSurfaceAllocSize: @(bytes), + (id)kIOSurfacePixelFormat: @0 + }); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) ptr[i] = fill; + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static IOSurfaceRef make_pattern_input_surface(int d_model, int seq) { + int count = d_model * seq; + IOSurfaceRef s = make_f32_surface(count, 0.0f); + IOSurfaceLock(s, 0, NULL); + float *ptr = (float *)IOSurfaceGetBaseAddress(s); + for (int c = 0; c < d_model; c++) { + for (int t = 0; t < seq; t++) { + ptr[c * seq + t] = sinf((float)(c * 0.001 + t * 0.01)) * 0.1f; + } + } + IOSurfaceUnlock(s, 0, NULL); + return s; +} + +static bool surface_all_finite(IOSurfaceRef s, int count, float *max_abs_out) { + bool ok = true; + float max_abs = 0.0f; + IOSurfaceLock(s, kIOSurfaceLockReadOnly, NULL); + const float *ptr = (const float *)IOSurfaceGetBaseAddress(s); + for (int i = 0; i < count; i++) { + float v = ptr[i]; + if (!isfinite(v)) ok = false; + float a = fabsf(v); + if (a > max_abs) max_abs = a; + } + IOSurfaceUnlock(s, kIOSurfaceLockReadOnly, NULL); + if (max_abs_out) *max_abs_out = max_abs; + return ok; +} + +static void add_blob(NSMutableDictionary *dict, NSString *mil_path, NSString *file_path) { + NSData *data = [NSData dataWithContentsOfFile:file_path]; + if (data) { + dict[mil_path] = @{@"offset": @0, @"data": data}; + } +} + +static NSDictionary *build_qproj_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_q_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_q_proj.bin"]); + return dict; +} + +static NSDictionary *build_kv_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/input_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"input_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_k_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_k_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/self_attn_v_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"self_attn_v_proj.bin"]); + return dict; +} + +static NSDictionary *build_ffn_wdict(int layer, NSString *blob_dir) { + NSMutableDictionary *dict = [NSMutableDictionary dictionary]; + NSString *prefix = [NSString stringWithFormat:@"%@/layer%d", blob_dir, layer]; + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/post_attention_layernorm.bin", layer], + [prefix stringByAppendingPathComponent:@"post_attention_layernorm.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_gate_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_gate_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_up_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_up_proj.bin"]); + add_blob(dict, [NSString stringWithFormat:@"@model_path/layer%d/mlp_down_proj.bin", layer], + [prefix stringByAppendingPathComponent:@"mlp_down_proj.bin"]); + return dict; +} + +static NSString *compile_graph(OrionGraph *g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString *mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static bool test_q_proj(NSString *blob_dir) { + int layer = 3; + int seq = 32; + int d = kQwen35_9B.d_model; + NSString *mil = compile_graph(orion_frontend_qwen35_prefill_q_proj(layer, seq, &kQwen35_9B)); + if (!mil || mil.length == 0) return false; + NSDictionary *wdict = build_qproj_wdict(layer, blob_dir); + OrionProgram *prog = orion_compile_mil(mil.UTF8String, wdict, "qwen35_9b_qproj"); + if (!prog) return false; + + IOSurfaceRef ioIn = make_pattern_input_surface(d, seq); + IOSurfaceRef ioOut = make_f32_surface(d * 2 * seq, 0.0f); + IOSurfaceRef ins[] = {ioIn}; + IOSurfaceRef outs[] = {ioOut}; + bool ok = orion_eval(prog, ins, 1, outs, 1); + float max_abs = 0.0f; + ok = ok && surface_all_finite(ioOut, d * 2 * seq, &max_abs) && max_abs > 1e-4f; + + CFRelease(ioIn); + CFRelease(ioOut); + orion_release_program(prog); + return ok; +} + +static bool test_kv_proj(NSString *blob_dir) { + int layer = 3; + int seq = 32; + int d = kQwen35_9B.d_model; + int kv = kQwen35_9B.n_kv_head * kQwen35_9B.head_dim; + NSString *mil = compile_graph(orion_frontend_qwen35_prefill_kv_proj(layer, seq, &kQwen35_9B)); + if (!mil || mil.length == 0) return false; + NSDictionary *wdict = build_kv_wdict(layer, blob_dir); + OrionProgram *prog = orion_compile_mil(mil.UTF8String, wdict, "qwen35_9b_kvproj"); + if (!prog) return false; + + IOSurfaceRef ioIn = make_pattern_input_surface(d, seq); + IOSurfaceRef ioK = make_f32_surface(kv * seq, 0.0f); + IOSurfaceRef ioV = make_f32_surface(kv * seq, 0.0f); + IOSurfaceRef ins[] = {ioIn}; + IOSurfaceRef outs[] = {ioK, ioV}; + bool ok = orion_eval(prog, ins, 1, outs, 2); + float max_k = 0.0f, max_v = 0.0f; + ok = ok && + surface_all_finite(ioK, kv * seq, &max_k) && + surface_all_finite(ioV, kv * seq, &max_v) && + max_k > 1e-4f && max_v > 1e-4f; + + CFRelease(ioIn); + CFRelease(ioK); + CFRelease(ioV); + orion_release_program(prog); + return ok; +} + +static bool test_ffn(NSString *blob_dir) { + int layer = 3; + int seq = 32; + int d = kQwen35_9B.d_model; + NSString *mil = compile_graph(orion_frontend_qwen35_prefill_ffn(layer, seq, &kQwen35_9B)); + if (!mil || mil.length == 0) return false; + NSDictionary *wdict = build_ffn_wdict(layer, blob_dir); + OrionProgram *prog = orion_compile_mil(mil.UTF8String, wdict, "qwen35_9b_ffn"); + if (!prog) return false; + + IOSurfaceRef ioIn = make_pattern_input_surface(d, seq); + IOSurfaceRef ioOut = make_f32_surface(d * seq, 0.0f); + IOSurfaceRef ins[] = {ioIn}; + IOSurfaceRef outs[] = {ioOut}; + bool ok = orion_eval(prog, ins, 1, outs, 1); + float max_abs = 0.0f; + ok = ok && surface_all_finite(ioOut, d * seq, &max_abs) && max_abs > 1e-4f; + + CFRelease(ioIn); + CFRelease(ioOut); + orion_release_program(prog); + return ok; +} + +int main(int argc, char **argv) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + if (!orion_ane_init()) { + fprintf(stderr, "orion_ane_init failed\n"); + return 3; + } + NSString *blob_dir = @(argv[1]); + printf("test_qwen35_9b_prefill_runtime_smoke:\n"); + TEST(q_proj); + TEST(kv_proj); + TEST(ffn); + printf("\n%d/%d passed\n", g_passed, g_passed + g_failed); + return g_failed > 0 ? 1 : 0; + } +} diff --git a/tests/test_qwen35_attention_shape_audit.m b/tests/test_qwen35_attention_shape_audit.m new file mode 100644 index 0000000..694b191 --- /dev/null +++ b/tests/test_qwen35_attention_shape_audit.m @@ -0,0 +1,135 @@ +#import +#import +#import "../model/weight_loader.h" + +static int product_of_shape(NSArray *shape) { + int total = 1; + for (id value in shape) { + if (![value respondsToSelector:@selector(intValue)]) { + return -1; + } + int dim = [value intValue]; + if (dim <= 0) { + return -1; + } + total *= dim; + } + return total; +} + +static NSDictionary* entry_for_path(NSArray *entries, NSString *relPath) { + for (id value in entries) { + if (![value isKindOfClass:[NSDictionary class]]) continue; + NSDictionary *entry = (NSDictionary *)value; + NSString *path = entry[@"path"]; + if ([path isEqualToString:relPath]) { + return entry; + } + } + return nil; +} + +static int validate_entry(NSArray *entries, NSString *blobDir, NSString *relPath) { + NSDictionary *entry = entry_for_path(entries, relPath); + if (!entry) { + fprintf(stderr, "FAIL: missing manifest entry for %s\n", relPath.UTF8String); + return 0; + } + + NSArray *shape = entry[@"shape"]; + if (![shape isKindOfClass:[NSArray class]] || [shape count] == 0) { + fprintf(stderr, "FAIL: missing shape for %s\n", relPath.UTF8String); + return 0; + } + + int expected = product_of_shape(shape); + if (expected <= 0) { + fprintf(stderr, "FAIL: invalid shape for %s\n", relPath.UTF8String); + return 0; + } + + NSString *absPath = [blobDir stringByAppendingPathComponent:relPath]; + int got = orion_blob_element_count(absPath.UTF8String); + if (got != expected) { + fprintf(stderr, "FAIL: %s element mismatch expected=%d got=%d\n", + relPath.UTF8String, expected, got); + return 0; + } + + printf(" %s shape=%s elements=%d\n", + relPath.UTF8String, + [[shape description] UTF8String], + got); + return 1; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + NSString *blobDir = [NSString stringWithUTF8String:argv[1]]; + NSString *manifestPath = [blobDir stringByAppendingPathComponent:@"manifest.json"]; + NSData *data = [NSData dataWithContentsOfFile:manifestPath]; + if (!data) { + fprintf(stderr, "FAIL: cannot read %s\n", manifestPath.UTF8String); + return 1; + } + + NSError *error = nil; + NSDictionary *manifest = [NSJSONSerialization JSONObjectWithData:data options:0 error:&error]; + if (![manifest isKindOfClass:[NSDictionary class]]) { + fprintf(stderr, "FAIL: invalid manifest json: %s\n", error.localizedDescription.UTF8String); + return 1; + } + + NSDictionary *runtime = manifest[@"runtime"]; + NSArray *entries = manifest[@"present_entries"]; + if (![runtime isKindOfClass:[NSDictionary class]] || ![entries isKindOfClass:[NSArray class]]) { + fprintf(stderr, "FAIL: manifest missing runtime/present_entries\n"); + return 1; + } + + int nLinear = 0; + int nFull = 0; + NSArray *layerTypes = runtime[@"layer_types"]; + if ([layerTypes isKindOfClass:[NSArray class]]) { + for (id value in layerTypes) { + if (![value isKindOfClass:[NSString class]]) continue; + NSString *layerType = (NSString *)value; + if ([layerType isEqualToString:@"linear_attention"]) nLinear += 1; + if ([layerType isEqualToString:@"full_attention"]) nFull += 1; + } + } + + printf("PASS: qwen35 attention shape audit\n"); + printf(" blob_dir=%s\n", blobDir.UTF8String); + printf(" n_linear_layers=%d\n", nLinear); + printf(" n_full_layers=%d\n", nFull); + printf(" attention_note=%s\n", "Qwen3.5 hybrid attention uses non-standard projected dimensions"); + + int ok = 1; + ok &= validate_entry(entries, blobDir, @"layer3/self_attn_q_proj.bin"); + ok &= validate_entry(entries, blobDir, @"layer3/self_attn_k_proj.bin"); + ok &= validate_entry(entries, blobDir, @"layer3/self_attn_v_proj.bin"); + ok &= validate_entry(entries, blobDir, @"layer3/self_attn_o_proj.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_in_proj_qkv.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_in_proj_z.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_in_proj_a.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_in_proj_b.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_out_proj.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_norm.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_dt_bias.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_a_log.bin"); + ok &= validate_entry(entries, blobDir, @"layer0/linear_attn_conv1d.bin"); + + if (!ok) { + return 1; + } + + printf(" next_blocker=%s\n", "derive full_attention and linear_attention semantics before CPU kernel implementation"); + return 0; + } +} diff --git a/tests/test_qwen35_cpu_forward_scaffold.m b/tests/test_qwen35_cpu_forward_scaffold.m new file mode 100644 index 0000000..e8d78ec --- /dev/null +++ b/tests/test_qwen35_cpu_forward_scaffold.m @@ -0,0 +1,123 @@ +#import +#import +#import +#import "../model/weight_loader.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +static int file_exists(const char *path) { + return [[NSFileManager defaultManager] fileExistsAtPath:[NSString stringWithUTF8String:path]]; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + float *embed_row = (float *)calloc((size_t)manifest->d_model, sizeof(float)); + float *final_norm = (float *)calloc((size_t)manifest->d_model, sizeof(float)); + float *input_layernorm = (float *)calloc((size_t)manifest->d_model, sizeof(float)); + if (!embed_row || !final_norm || !input_layernorm) { + fprintf(stderr, "FAIL: allocation failed\n"); + free(embed_row); + free(final_norm); + free(input_layernorm); + orion_qwen35_manifest_free(manifest); + return 1; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, manifest->d_model, embed_row)) { + fprintf(stderr, "FAIL: failed to read embedding row\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/model/final_norm.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, manifest->d_model, final_norm)) { + fprintf(stderr, "FAIL: failed to read final norm\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer0/input_layernorm.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, manifest->d_model, input_layernorm)) { + fprintf(stderr, "FAIL: failed to read layer0 input layernorm\n"); + goto fail; + } + + int detected_linear = 0; + int detected_full = 0; + int first_linear_idx = -1; + int first_full_idx = -1; + for (int i = 0; i < manifest->n_layer; i++) { + char linear_qkv[2048]; + char full_q[2048]; + snprintf(linear_qkv, sizeof(linear_qkv), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, i); + snprintf(full_q, sizeof(full_q), "%s/layer%d/self_attn_q_proj.bin", blob_dir, i); + if (file_exists(linear_qkv)) { + detected_linear += 1; + if (first_linear_idx < 0) first_linear_idx = i; + } else if (file_exists(full_q)) { + detected_full += 1; + if (first_full_idx < 0) first_full_idx = i; + } else { + fprintf(stderr, "FAIL: layer%d has neither linear nor full attention blobs\n", i); + goto fail; + } + } + + double embed_abs = abs_sum(embed_row, manifest->d_model); + double final_norm_abs = abs_sum(final_norm, manifest->d_model); + double input_ln_abs = abs_sum(input_layernorm, manifest->d_model); + + if (embed_abs <= 0.0 || final_norm_abs <= 0.0 || input_ln_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero tensor content\n"); + goto fail; + } + if (detected_linear != manifest->n_linear_layers || detected_full != manifest->n_full_layers) { + fprintf(stderr, "FAIL: layer topology mismatch (detected=%d/%d manifest=%d/%d)\n", + detected_linear, detected_full, manifest->n_linear_layers, manifest->n_full_layers); + goto fail; + } + + printf("PASS: qwen35 cpu forward scaffold\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" token_id=%d\n", 0); + printf(" embed_abs_sum=%.6f\n", embed_abs); + printf(" final_norm_abs_sum=%.6f\n", final_norm_abs); + printf(" layer0_input_layernorm_abs_sum=%.6f\n", input_ln_abs); + printf(" detected_linear_layers=%d\n", detected_linear); + printf(" detected_full_layers=%d\n", detected_full); + printf(" first_linear_layer_idx=%d\n", first_linear_idx); + printf(" first_full_layer_idx=%d\n", first_full_idx); + printf(" next_blocker=%s\n", "implement linear_attention/full_attention CPU kernels"); + + free(embed_row); + free(final_norm); + free(input_layernorm); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(embed_row); + free(final_norm); + free(input_layernorm); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_decode_loop_cpu_smoke.m b/tests/test_qwen35_decode_loop_cpu_smoke.m new file mode 100644 index 0000000..5ca0c23 --- /dev/null +++ b/tests/test_qwen35_decode_loop_cpu_smoke.m @@ -0,0 +1,282 @@ +#import +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../tokenizer/gpt2_bpe.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) total += fabs((double)x[i]); + return total; +} + +static NSDictionary* load_json(NSString* path) { + NSData* data = [NSData dataWithContentsOfFile:path]; + if (!data) return nil; + return [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +static int argmax_vocab_from_tied_embedding(const char *blob_dir, + const float *hidden, + int d_model, + int vocab, + int *top_id, + float *top_logit) { + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + + float *row = (float *)malloc((size_t)d_model * sizeof(float)); + if (!row) return 0; + + int best_id = -1; + float best_logit = -INFINITY; + for (int tok = 0; tok < vocab; tok++) { + if (!orion_read_blob_row_f32(path, tok, d_model, row)) { + free(row); + return 0; + } + float dot = 0.0f; + for (int i = 0; i < d_model; i++) dot += hidden[i] * row[i]; + if (dot > best_logit) { + best_logit = dot; + best_id = tok; + } + } + + free(row); + *top_id = best_id; + *top_logit = best_logit; + return 1; +} + +static int compute_top1_from_sequence(const char *blob_dir, + OrionQwen35Manifest *manifest, + const int *token_ids, + int seq_len, + int *top_id, + float *top_logit, + double *last_hidden_abs_sum) { + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int num_k_heads = 16; + const int num_v_heads = 16; + const int head_k_dim = 128; + const int head_v_dim = 128; + const int value_dim = num_v_heads * head_v_dim; + const int conv_kernel = 4; + + float *hidden = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *last_hidden = (float *)calloc((size_t)d_model, sizeof(float)); + float *final_norm = NULL; + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch || !last_hidden) goto fail; + + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + for (int s = 0; s < seq_len; s++) { + if (!orion_read_blob_row_f32(embed_path, token_ids[s], d_model, hidden + s * d_model)) goto fail; + } + + for (int layer_idx = 0; layer_idx < manifest->n_layer; layer_idx++) { + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + memset(mixer_out, 0, (size_t)seq_len * d_model * sizeof(float)); + + int is_full_layer = ((layer_idx + 1) % 4 == 0); + if (is_full_layer) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", (value_dim + value_dim + value_dim) * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", (value_dim + value_dim + value_dim) * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) hidden[i] += mixer_out[i]; + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + for (int i = 0; i < seq_len * d_model; i++) hidden[i] += mlp_out[i]; + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + char final_norm_path[2048]; + snprintf(final_norm_path, sizeof(final_norm_path), "%s/model/final_norm.bin", blob_dir); + final_norm = orion_read_blob_f32_exact(final_norm_path, d_model); + if (!final_norm) goto fail; + + orion_qwen_cpu_rmsnorm(hidden + (seq_len - 1) * d_model, final_norm, d_model, 1e-6f, last_hidden); + if (!argmax_vocab_from_tied_embedding(blob_dir, last_hidden, d_model, manifest->vocab, top_id, top_logit)) goto fail; + *last_hidden_abs_sum = abs_sum(last_hidden, d_model); + + free(hidden); free(normed); free(mixer_out); free(mlp_out); free(scratch); free(last_hidden); free(final_norm); + return 1; + +fail: + free(hidden); free(normed); free(mixer_out); free(mlp_out); free(scratch); free(last_hidden); free(final_norm); + return 0; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + NSString* tokDir = [NSString stringWithUTF8String:argv[2]]; + NSDictionary* meta = load_json([tokDir stringByAppendingPathComponent:@"meta.json"]); + if (!meta) { + fprintf(stderr, "FAIL: missing tokenizer meta.json\n"); + return 1; + } + NSString* regex = meta[@"regex_pattern"]; + NSString* vocabPath = [tokDir stringByAppendingPathComponent:@"vocab.json"]; + NSString* mergesPath = [tokDir stringByAppendingPathComponent:@"merges.txt"]; + OrionGPT2Tokenizer* tok = orion_gpt2_tokenizer_load_with_regex(vocabPath.UTF8String, mergesPath.UTF8String, regex.UTF8String); + if (!tok) { + fprintf(stderr, "FAIL: tokenizer load failed\n"); + return 1; + } + + OrionQwen35Manifest* manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + const char *prompt = "YES"; + int tokens[64] = {0}; + int prompt_len = orion_gpt2_encode(tok, prompt, tokens, 64); + if (prompt_len <= 0) { + fprintf(stderr, "FAIL: prompt encode failed\n"); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + int gen_ids[2] = {0}; + float gen_logits[2] = {0.0f, 0.0f}; + double final_hidden_abs = 0.0; + int seq[128] = {0}; + memcpy(seq, tokens, (size_t)prompt_len * sizeof(int)); + int seq_len = prompt_len; + + for (int step = 0; step < 2; step++) { + int next_id = -1; + float next_logit = -INFINITY; + if (!compute_top1_from_sequence(blob_dir, manifest, seq, seq_len, &next_id, &next_logit, &final_hidden_abs)) { + fprintf(stderr, "FAIL: top1 compute failed at step %d\n", step); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + gen_ids[step] = next_id; + gen_logits[step] = next_logit; + seq[seq_len++] = next_id; + } + + char* decoded_prompt = orion_gpt2_decode(tok, tokens, prompt_len); + char* decoded_generated = orion_gpt2_decode(tok, gen_ids, 2); + char* decoded_full = orion_gpt2_decode(tok, seq, seq_len); + + if (!decoded_prompt || !decoded_generated || !decoded_full) { + fprintf(stderr, "FAIL: decode returned NULL\n"); + free(decoded_prompt); free(decoded_generated); free(decoded_full); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 1; + } + + printf("PASS: qwen35 decode loop cpu smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" tokenizer_dir=%s\n", tokDir.UTF8String); + printf(" prompt=%s\n", prompt); + printf(" prompt_len=%d\n", prompt_len); + printf(" gen_len=2\n"); + printf(" gen_token_0=%d\n", gen_ids[0]); + printf(" gen_token_1=%d\n", gen_ids[1]); + printf(" gen_logit_0=%.6f\n", gen_logits[0]); + printf(" gen_logit_1=%.6f\n", gen_logits[1]); + printf(" final_hidden_abs_sum=%.6f\n", final_hidden_abs); + printf(" decoded_prompt=%s\n", decoded_prompt); + printf(" decoded_generated=%s\n", decoded_generated); + printf(" decoded_full=%s\n", decoded_full); + + free(decoded_prompt); + free(decoded_generated); + free(decoded_full); + orion_qwen35_manifest_free(manifest); + orion_gpt2_tokenizer_free(tok); + return 0; + } +} diff --git a/tests/test_qwen35_full_attention_cpu_smoke.m b/tests/test_qwen35_full_attention_cpu_smoke.m new file mode 100644 index 0000000..a0f9581 --- /dev/null +++ b/tests/test_qwen35_full_attention_cpu_smoke.m @@ -0,0 +1,163 @@ +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int layer_idx = 3; + const int seq_len = 2; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *x_seq = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *attn_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *residual = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *post_ln_weight = (float *)calloc((size_t)d_model, sizeof(float)); + float *rms_buf = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!x_seq || !attn_out || !residual || !post_ln_weight || !rms_buf || !mlp_out) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, d_model, x_seq + 0 * d_model) || + !orion_read_blob_row_f32(path, 1, d_model, x_seq + 1 * d_model)) { + fprintf(stderr, "FAIL: failed to read token embeddings\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer%d/input_layernorm.bin", blob_dir, layer_idx); + float *input_ln = orion_read_blob_f32_exact(path, d_model); + snprintf(path, sizeof(path), "%s/layer%d/post_attention_layernorm.bin", blob_dir, layer_idx); + float *post_ln = orion_read_blob_f32_exact(path, d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + float *q_proj = orion_read_blob_f32_exact(path, (q_dim * 2) * d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_k_proj.bin", blob_dir, layer_idx); + float *k_proj = orion_read_blob_f32_exact(path, kv_dim * d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_v_proj.bin", blob_dir, layer_idx); + float *v_proj = orion_read_blob_f32_exact(path, kv_dim * d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_o_proj.bin", blob_dir, layer_idx); + float *o_proj = orion_read_blob_f32_exact(path, d_model * q_dim); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_q_norm.bin", blob_dir, layer_idx); + float *q_norm = orion_read_blob_f32_exact(path, head_dim); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_k_norm.bin", blob_dir, layer_idx); + float *k_norm = orion_read_blob_f32_exact(path, head_dim); + snprintf(path, sizeof(path), "%s/layer%d/mlp_gate_proj.bin", blob_dir, layer_idx); + float *gate_proj = orion_read_blob_f32_exact(path, d_ff * d_model); + snprintf(path, sizeof(path), "%s/layer%d/mlp_up_proj.bin", blob_dir, layer_idx); + float *up_proj = orion_read_blob_f32_exact(path, d_ff * d_model); + snprintf(path, sizeof(path), "%s/layer%d/mlp_down_proj.bin", blob_dir, layer_idx); + float *down_proj = orion_read_blob_f32_exact(path, d_model * d_ff); + + if (!input_ln || !post_ln || !q_proj || !k_proj || !v_proj || !o_proj || + !q_norm || !k_norm || !gate_proj || !up_proj || !down_proj) { + fprintf(stderr, "FAIL: failed to read one or more full-attention layer tensors\n"); + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); + free(q_norm); free(k_norm); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(x_seq + s * d_model, input_ln, d_model, 1e-6f, rms_buf + s * d_model); + } + + orion_qwen_cpu_full_attention_prefill_no_rope( + rms_buf, + seq_len, + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + d_model, + n_head, + n_kv_head, + head_dim, + attn_out + ); + + for (int i = 0; i < seq_len * d_model; i++) { + residual[i] = x_seq[i] + attn_out[i]; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(residual + s * d_model, post_ln, d_model, 1e-6f, rms_buf + s * d_model); + orion_qwen_cpu_swiglu_ffn(rms_buf + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + + double input_abs = abs_sum(x_seq, seq_len * d_model); + double attn_abs = abs_sum(attn_out, seq_len * d_model); + double residual_abs = abs_sum(residual, seq_len * d_model); + double mlp_abs = abs_sum(mlp_out, seq_len * d_model); + + free(input_ln); free(post_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); + free(q_norm); free(k_norm); free(gate_proj); free(up_proj); free(down_proj); + + if (input_abs <= 0.0 || attn_abs <= 0.0 || residual_abs <= 0.0 || mlp_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero attention/MLP activations\n"); + goto fail; + } + + printf("PASS: qwen35 full attention cpu smoke no rope\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" layer_idx=%d\n", layer_idx); + printf(" seq_len=%d\n", seq_len); + printf(" q_dim=%d\n", q_dim); + printf(" kv_dim=%d\n", kv_dim); + printf(" input_abs_sum=%.6f\n", input_abs); + printf(" attn_abs_sum=%.6f\n", attn_abs); + printf(" residual_abs_sum=%.6f\n", residual_abs); + printf(" mlp_abs_sum=%.6f\n", mlp_abs); + printf(" next_blocker=%s\n", "implement rotary embeddings and linear_attention CPU kernels"); + + free(x_seq); + free(attn_out); + free(residual); + free(post_ln_weight); + free(rms_buf); + free(mlp_out); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(x_seq); + free(attn_out); + free(residual); + free(post_ln_weight); + free(rms_buf); + free(mlp_out); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_full_attention_rope_cpu_smoke.m b/tests/test_qwen35_full_attention_rope_cpu_smoke.m new file mode 100644 index 0000000..dcead68 --- /dev/null +++ b/tests/test_qwen35_full_attention_rope_cpu_smoke.m @@ -0,0 +1,132 @@ +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int layer_idx = 3; + const int seq_len = 2; + const int d_model = manifest->d_model; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + + float *x_seq = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *rms_buf = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *attn_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!x_seq || !rms_buf || !attn_out) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, d_model, x_seq + 0 * d_model) || + !orion_read_blob_row_f32(path, 1, d_model, x_seq + 1 * d_model)) { + fprintf(stderr, "FAIL: failed to read token embeddings\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer%d/input_layernorm.bin", blob_dir, layer_idx); + float *input_ln = orion_read_blob_f32_exact(path, d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_q_proj.bin", blob_dir, layer_idx); + float *q_proj = orion_read_blob_f32_exact(path, (q_dim * 2) * d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_k_proj.bin", blob_dir, layer_idx); + float *k_proj = orion_read_blob_f32_exact(path, kv_dim * d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_v_proj.bin", blob_dir, layer_idx); + float *v_proj = orion_read_blob_f32_exact(path, kv_dim * d_model); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_o_proj.bin", blob_dir, layer_idx); + float *o_proj = orion_read_blob_f32_exact(path, d_model * q_dim); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_q_norm.bin", blob_dir, layer_idx); + float *q_norm = orion_read_blob_f32_exact(path, head_dim); + snprintf(path, sizeof(path), "%s/layer%d/self_attn_k_norm.bin", blob_dir, layer_idx); + float *k_norm = orion_read_blob_f32_exact(path, head_dim); + + if (!input_ln || !q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + fprintf(stderr, "FAIL: failed to read one or more full-attention rope tensors\n"); + free(input_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(x_seq + s * d_model, input_ln, d_model, 1e-6f, rms_buf + s * d_model); + } + + orion_qwen_cpu_full_attention_prefill_with_rope( + rms_buf, + seq_len, + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + d_model, + n_head, + n_kv_head, + head_dim, + manifest->rope_theta, + manifest->partial_rotary_factor, + attn_out + ); + + double input_abs = abs_sum(x_seq, seq_len * d_model); + double attn_abs = abs_sum(attn_out, seq_len * d_model); + + free(input_ln); free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + + if (input_abs <= 0.0 || attn_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero rope attention activations\n"); + goto fail; + } + + printf("PASS: qwen35 full attention cpu smoke with rope\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" layer_idx=%d\n", layer_idx); + printf(" seq_len=%d\n", seq_len); + printf(" rope_theta=%.1f\n", manifest->rope_theta); + printf(" partial_rotary_factor=%.4f\n", manifest->partial_rotary_factor); + printf(" rotary_dim=%d\n", manifest->rotary_dim); + printf(" input_abs_sum=%.6f\n", input_abs); + printf(" attn_abs_sum=%.6f\n", attn_abs); + printf(" next_blocker=%s\n", "implement hybrid decoder-layer dispatch"); + + free(x_seq); + free(rms_buf); + free(attn_out); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(x_seq); + free(rms_buf); + free(attn_out); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_hybrid_dispatch_cpu_smoke.m b/tests/test_qwen35_hybrid_dispatch_cpu_smoke.m new file mode 100644 index 0000000..fb383ab --- /dev/null +++ b/tests/test_qwen35_hybrid_dispatch_cpu_smoke.m @@ -0,0 +1,185 @@ +#import +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +static float *load_exact(const char *blob_dir, int layer_idx, const char *suffix, int count) { + char path[2048]; + snprintf(path, sizeof(path), "%s/layer%d/%s", blob_dir, layer_idx, suffix); + return orion_read_blob_f32_exact(path, count); +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int seq_len = 2; + const int d_model = manifest->d_model; + const int d_ff = manifest->d_ff; + const int n_head = manifest->n_head; + const int n_kv_head = manifest->n_kv_head; + const int head_dim = manifest->head_dim; + const int q_dim = n_head * head_dim; + const int kv_dim = n_kv_head * head_dim; + const int num_k_heads = 16; + const int num_v_heads = 16; + const int head_k_dim = 128; + const int head_v_dim = 128; + const int value_dim = num_v_heads * head_v_dim; + const int conv_kernel = 4; + + float *hidden = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mixer_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *scratch = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!hidden || !normed || !mixer_out || !mlp_out || !scratch) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char embed_path[2048]; + snprintf(embed_path, sizeof(embed_path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(embed_path, 0, d_model, hidden + 0 * d_model) || + !orion_read_blob_row_f32(embed_path, 1, d_model, hidden + 1 * d_model)) { + fprintf(stderr, "FAIL: failed to read token embeddings\n"); + goto fail; + } + + const int layer_indices[4] = {0, 1, 2, 3}; + for (int idx = 0; idx < 4; idx++) { + int layer_idx = layer_indices[idx]; + + float *input_ln = load_exact(blob_dir, layer_idx, "input_layernorm.bin", d_model); + float *post_ln = load_exact(blob_dir, layer_idx, "post_attention_layernorm.bin", d_model); + float *gate_proj = load_exact(blob_dir, layer_idx, "mlp_gate_proj.bin", d_ff * d_model); + float *up_proj = load_exact(blob_dir, layer_idx, "mlp_up_proj.bin", d_ff * d_model); + float *down_proj = load_exact(blob_dir, layer_idx, "mlp_down_proj.bin", d_model * d_ff); + if (!input_ln || !post_ln || !gate_proj || !up_proj || !down_proj) { + fprintf(stderr, "FAIL: missing layer %d norm/mlp tensors\n", layer_idx); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, input_ln, d_model, 1e-6f, normed + s * d_model); + } + + memset(mixer_out, 0, (size_t)seq_len * d_model * sizeof(float)); + + if (layer_idx == 3) { + float *q_proj = load_exact(blob_dir, layer_idx, "self_attn_q_proj.bin", (q_dim * 2) * d_model); + float *k_proj = load_exact(blob_dir, layer_idx, "self_attn_k_proj.bin", kv_dim * d_model); + float *v_proj = load_exact(blob_dir, layer_idx, "self_attn_v_proj.bin", kv_dim * d_model); + float *o_proj = load_exact(blob_dir, layer_idx, "self_attn_o_proj.bin", d_model * q_dim); + float *q_norm = load_exact(blob_dir, layer_idx, "self_attn_q_norm.bin", head_dim); + float *k_norm = load_exact(blob_dir, layer_idx, "self_attn_k_norm.bin", head_dim); + if (!q_proj || !k_proj || !v_proj || !o_proj || !q_norm || !k_norm) { + fprintf(stderr, "FAIL: missing full-attention tensors for layer %d\n", layer_idx); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_full_attention_prefill_with_rope( + normed, seq_len, q_proj, k_proj, v_proj, o_proj, q_norm, k_norm, + d_model, n_head, n_kv_head, head_dim, manifest->rope_theta, manifest->partial_rotary_factor, + mixer_out + ); + free(q_proj); free(k_proj); free(v_proj); free(o_proj); free(q_norm); free(k_norm); + } else { + float *in_proj_qkv = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_qkv.bin", (value_dim + value_dim + value_dim) * d_model); + float *in_proj_z = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_z.bin", value_dim * d_model); + float *in_proj_a = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_a.bin", num_v_heads * d_model); + float *in_proj_b = load_exact(blob_dir, layer_idx, "linear_attn_in_proj_b.bin", num_v_heads * d_model); + float *conv1d = load_exact(blob_dir, layer_idx, "linear_attn_conv1d.bin", (value_dim + value_dim + value_dim) * conv_kernel); + float *dt_bias = load_exact(blob_dir, layer_idx, "linear_attn_dt_bias.bin", num_v_heads); + float *a_log = load_exact(blob_dir, layer_idx, "linear_attn_a_log.bin", num_v_heads); + float *norm_weight = load_exact(blob_dir, layer_idx, "linear_attn_norm.bin", head_v_dim); + float *out_proj = load_exact(blob_dir, layer_idx, "linear_attn_out_proj.bin", d_model * value_dim); + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + fprintf(stderr, "FAIL: missing linear-attention tensors for layer %d\n", layer_idx); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + goto fail; + } + orion_qwen_cpu_linear_attention_recurrent_prefill( + normed, seq_len, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + mixer_out + ); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + } + + for (int i = 0; i < seq_len * d_model; i++) { + hidden[i] += mixer_out[i]; + } + + for (int s = 0; s < seq_len; s++) { + orion_qwen_cpu_rmsnorm(hidden + s * d_model, post_ln, d_model, 1e-6f, scratch + s * d_model); + orion_qwen_cpu_swiglu_ffn(scratch + s * d_model, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out + s * d_model); + } + + for (int i = 0; i < seq_len * d_model; i++) { + hidden[i] += mlp_out[i]; + } + + free(input_ln); free(post_ln); free(gate_proj); free(up_proj); free(down_proj); + } + + double hidden_abs = abs_sum(hidden, seq_len * d_model); + if (hidden_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero hybrid dispatch output\n"); + goto fail; + } + + printf("PASS: qwen35 hybrid dispatch cpu smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" seq_len=%d\n", seq_len); + printf(" layers_run=%d\n", 4); + printf(" last_full_layer_idx=%d\n", 3); + printf(" hidden_abs_sum=%.6f\n", hidden_abs); + printf(" next_blocker=%s\n", "implement full 24-layer 1-token logits smoke"); + + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(hidden); + free(normed); + free(mixer_out); + free(mlp_out); + free(scratch); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_infer_prep.m b/tests/test_qwen35_infer_prep.m new file mode 100644 index 0000000..0b16a33 --- /dev/null +++ b/tests/test_qwen35_infer_prep.m @@ -0,0 +1,75 @@ +#import +#import +#import "../core/model_registry.h" +#import "../model/weight_loader.h" + +static int require_match(const char *label, int lhs, int rhs) { + if (lhs != rhs) { + fprintf(stderr, "FAIL: %s mismatch (registry=%d manifest=%d)\n", label, lhs, rhs); + return 0; + } + return 1; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const OrionModelSpec *spec = orion_model_lookup("qwen35_08b"); + if (!spec) { + fprintf(stderr, "FAIL: registry lookup returned NULL\n"); + return 1; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(argv[1]); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + printf("PASS: qwen35 cpu infer prep\n"); + printf(" model_name=%s\n", spec->name); + printf(" blob_dir=%s\n", manifest->blob_dir); + printf(" manifest_path=%s\n", manifest->manifest_path); + printf(" registry_n_layer=%d\n", spec->config.n_layer); + printf(" manifest_n_layer=%d\n", manifest->n_layer); + printf(" registry_d_model=%d\n", spec->config.d_model); + printf(" manifest_d_model=%d\n", manifest->d_model); + printf(" registry_d_ff=%d\n", spec->config.hidden_dim); + printf(" manifest_d_ff=%d\n", manifest->d_ff); + printf(" registry_n_head=%d\n", spec->config.n_head); + printf(" manifest_n_head=%d\n", manifest->n_head); + printf(" registry_n_kv_head=%d\n", spec->config.n_kv_head); + printf(" manifest_n_kv_head=%d\n", manifest->n_kv_head); + printf(" registry_head_dim=%d\n", spec->config.head_dim); + printf(" manifest_head_dim=%d\n", manifest->head_dim); + printf(" registry_vocab=%d\n", spec->config.vocab); + printf(" manifest_vocab=%d\n", manifest->vocab); + printf(" registry_max_seq=%d\n", spec->config.max_seq); + printf(" manifest_max_seq=%d\n", manifest->max_seq); + printf(" manifest_n_linear_layers=%d\n", manifest->n_linear_layers); + printf(" manifest_n_full_layers=%d\n", manifest->n_full_layers); + + int ok = 1; + ok &= require_match("n_layer", spec->config.n_layer, manifest->n_layer); + ok &= require_match("d_model", spec->config.d_model, manifest->d_model); + ok &= require_match("d_ff", spec->config.hidden_dim, manifest->d_ff); + ok &= require_match("n_head", spec->config.n_head, manifest->n_head); + ok &= require_match("n_kv_head", spec->config.n_kv_head, manifest->n_kv_head); + ok &= require_match("head_dim", spec->config.head_dim, manifest->head_dim); + ok &= require_match("vocab", spec->config.vocab, manifest->vocab); + ok &= require_match("max_seq", spec->config.max_seq, manifest->max_seq); + + if (manifest->n_linear_layers != 18 || manifest->n_full_layers != 6) { + fprintf(stderr, "FAIL: unexpected qwen35 hybrid layer topology (%d linear / %d full)\n", + manifest->n_linear_layers, manifest->n_full_layers); + ok = 0; + } + + orion_qwen35_manifest_free(manifest); + return ok ? 0 : 1; + } +} diff --git a/tests/test_qwen35_linear_attention_prep_cpu_smoke.m b/tests/test_qwen35_linear_attention_prep_cpu_smoke.m new file mode 100644 index 0000000..2b2ebb9 --- /dev/null +++ b/tests/test_qwen35_linear_attention_prep_cpu_smoke.m @@ -0,0 +1,157 @@ +#import +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int layer_idx = 0; + const int seq_len = 2; + const int d_model = manifest->d_model; + const int num_k_heads = 16; + const int num_v_heads = 16; + const int head_k_dim = 128; + const int head_v_dim = 128; + const int conv_kernel = 4; + const int key_dim = num_k_heads * head_k_dim; + const int value_dim = num_v_heads * head_v_dim; + + float *x_seq = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *query = (float *)calloc((size_t)seq_len * key_dim, sizeof(float)); + float *key = (float *)calloc((size_t)seq_len * key_dim, sizeof(float)); + float *value = (float *)calloc((size_t)seq_len * value_dim, sizeof(float)); + float *z = (float *)calloc((size_t)seq_len * value_dim, sizeof(float)); + float *beta = (float *)calloc((size_t)seq_len * num_v_heads, sizeof(float)); + float *g = (float *)calloc((size_t)seq_len * num_v_heads, sizeof(float)); + if (!x_seq || !query || !key || !value || !z || !beta || !g) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, d_model, x_seq + 0 * d_model) || + !orion_read_blob_row_f32(path, 1, d_model, x_seq + 1 * d_model)) { + fprintf(stderr, "FAIL: failed to read token embeddings\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + float *in_proj_qkv = orion_read_blob_f32_exact(path, (key_dim * 2 + value_dim) * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_z.bin", blob_dir, layer_idx); + float *in_proj_z = orion_read_blob_f32_exact(path, value_dim * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_a.bin", blob_dir, layer_idx); + float *in_proj_a = orion_read_blob_f32_exact(path, num_v_heads * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_b.bin", blob_dir, layer_idx); + float *in_proj_b = orion_read_blob_f32_exact(path, num_v_heads * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + float *conv1d = orion_read_blob_f32_exact(path, (key_dim * 2 + value_dim) * conv_kernel); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + float *dt_bias = orion_read_blob_f32_exact(path, num_v_heads); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_a_log.bin", blob_dir, layer_idx); + float *a_log = orion_read_blob_f32_exact(path, num_v_heads); + + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || !dt_bias || !a_log) { + fprintf(stderr, "FAIL: failed to read one or more linear-attention tensors\n"); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); + goto fail; + } + + orion_qwen_cpu_linear_attention_prep( + x_seq, + seq_len, + in_proj_qkv, + in_proj_z, + in_proj_a, + in_proj_b, + conv1d, + dt_bias, + a_log, + d_model, + num_k_heads, + num_v_heads, + head_k_dim, + head_v_dim, + conv_kernel, + query, + key, + value, + z, + beta, + g + ); + + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); + + double q_abs = abs_sum(query, seq_len * key_dim); + double k_abs = abs_sum(key, seq_len * key_dim); + double v_abs = abs_sum(value, seq_len * value_dim); + double z_abs = abs_sum(z, seq_len * value_dim); + double beta_abs = abs_sum(beta, seq_len * num_v_heads); + double g_abs = abs_sum(g, seq_len * num_v_heads); + + if (q_abs <= 0.0 || k_abs <= 0.0 || v_abs <= 0.0 || z_abs <= 0.0 || beta_abs <= 0.0 || g_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero linear attention prep activations\n"); + goto fail; + } + + printf("PASS: qwen35 linear attention prep cpu smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" layer_idx=%d\n", layer_idx); + printf(" seq_len=%d\n", seq_len); + printf(" key_dim=%d\n", key_dim); + printf(" value_dim=%d\n", value_dim); + printf(" query_abs_sum=%.6f\n", q_abs); + printf(" key_abs_sum=%.6f\n", k_abs); + printf(" value_abs_sum=%.6f\n", v_abs); + printf(" z_abs_sum=%.6f\n", z_abs); + printf(" beta_abs_sum=%.6f\n", beta_abs); + printf(" g_abs_sum=%.6f\n", g_abs); + printf(" next_blocker=%s\n", "implement gated delta recurrent core and hybrid attention dispatch"); + + free(x_seq); + free(query); + free(key); + free(value); + free(z); + free(beta); + free(g); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(x_seq); + free(query); + free(key); + free(value); + free(z); + free(beta); + free(g); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_linear_attention_recurrent_cpu_smoke.m b/tests/test_qwen35_linear_attention_recurrent_cpu_smoke.m new file mode 100644 index 0000000..6be6c92 --- /dev/null +++ b/tests/test_qwen35_linear_attention_recurrent_cpu_smoke.m @@ -0,0 +1,116 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + const int layer_idx = 0; + const int seq_len = 2; + const int d_model = manifest->d_model; + const int num_k_heads = 16; + const int num_v_heads = 16; + const int head_k_dim = 128; + const int head_v_dim = 128; + const int value_dim = num_v_heads * head_v_dim; + const int conv_kernel = 4; + + float *x_seq = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + float *out_seq = (float *)calloc((size_t)seq_len * d_model, sizeof(float)); + if (!x_seq || !out_seq) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, d_model, x_seq + 0 * d_model) || + !orion_read_blob_row_f32(path, 1, d_model, x_seq + 1 * d_model)) { + fprintf(stderr, "FAIL: failed to read token embeddings\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_qkv.bin", blob_dir, layer_idx); + float *in_proj_qkv = orion_read_blob_f32_exact(path, (value_dim + value_dim + value_dim) * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_z.bin", blob_dir, layer_idx); + float *in_proj_z = orion_read_blob_f32_exact(path, value_dim * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_a.bin", blob_dir, layer_idx); + float *in_proj_a = orion_read_blob_f32_exact(path, num_v_heads * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_in_proj_b.bin", blob_dir, layer_idx); + float *in_proj_b = orion_read_blob_f32_exact(path, num_v_heads * d_model); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_conv1d.bin", blob_dir, layer_idx); + float *conv1d = orion_read_blob_f32_exact(path, (value_dim + value_dim + value_dim) * conv_kernel); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_dt_bias.bin", blob_dir, layer_idx); + float *dt_bias = orion_read_blob_f32_exact(path, num_v_heads); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_a_log.bin", blob_dir, layer_idx); + float *a_log = orion_read_blob_f32_exact(path, num_v_heads); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_norm.bin", blob_dir, layer_idx); + float *norm_weight = orion_read_blob_f32_exact(path, head_v_dim); + snprintf(path, sizeof(path), "%s/layer%d/linear_attn_out_proj.bin", blob_dir, layer_idx); + float *out_proj = orion_read_blob_f32_exact(path, d_model * value_dim); + + if (!in_proj_qkv || !in_proj_z || !in_proj_a || !in_proj_b || !conv1d || + !dt_bias || !a_log || !norm_weight || !out_proj) { + fprintf(stderr, "FAIL: failed to read one or more linear-attention recurrent tensors\n"); + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + goto fail; + } + + orion_qwen_cpu_linear_attention_recurrent_prefill( + x_seq, seq_len, + in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, dt_bias, a_log, norm_weight, out_proj, + d_model, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel, + out_seq + ); + + free(in_proj_qkv); free(in_proj_z); free(in_proj_a); free(in_proj_b); + free(conv1d); free(dt_bias); free(a_log); free(norm_weight); free(out_proj); + + double out_abs = abs_sum(out_seq, seq_len * d_model); + if (out_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero recurrent linear attention output\n"); + goto fail; + } + + printf("PASS: qwen35 linear attention recurrent cpu smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" layer_idx=%d\n", layer_idx); + printf(" seq_len=%d\n", seq_len); + printf(" out_abs_sum=%.6f\n", out_abs); + printf(" next_blocker=%s\n", "implement rotary embeddings and hybrid 1-token logits dispatch"); + + free(x_seq); + free(out_seq); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(x_seq); + free(out_seq); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_manifest_loader.m b/tests/test_qwen35_manifest_loader.m new file mode 100644 index 0000000..837b614 --- /dev/null +++ b/tests/test_qwen35_manifest_loader.m @@ -0,0 +1,54 @@ +#import +#import +#import "../model/weight_loader.h" + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(argv[1]); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + printf("PASS: qwen35 manifest loader\n"); + printf(" blob_dir=%s\n", manifest->blob_dir); + printf(" manifest_path=%s\n", manifest->manifest_path); + printf(" n_layer=%d\n", manifest->n_layer); + printf(" d_model=%d\n", manifest->d_model); + printf(" d_ff=%d\n", manifest->d_ff); + printf(" n_head=%d\n", manifest->n_head); + printf(" n_kv_head=%d\n", manifest->n_kv_head); + printf(" head_dim=%d\n", manifest->head_dim); + printf(" vocab=%d\n", manifest->vocab); + printf(" max_seq=%d\n", manifest->max_seq); + printf(" tie_word_embeddings=%d\n", manifest->tie_word_embeddings); + printf(" n_entries=%d\n", manifest->n_entries); + printf(" n_alias_entries=%d\n", manifest->n_alias_entries); + printf(" n_linear_layers=%d\n", manifest->n_linear_layers); + printf(" n_full_layers=%d\n", manifest->n_full_layers); + + if (manifest->n_layer != 24 || manifest->n_linear_layers != 18 || manifest->n_full_layers != 6) { + fprintf(stderr, "FAIL: unexpected layer topology\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (manifest->d_model != 1024 || manifest->d_ff != 3584 || manifest->vocab != 248320) { + fprintf(stderr, "FAIL: unexpected runtime dimensions\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + if (manifest->n_alias_entries < 1 || manifest->tie_word_embeddings != 1) { + fprintf(stderr, "FAIL: expected tied-embedding alias metadata\n"); + orion_qwen35_manifest_free(manifest); + return 1; + } + + orion_qwen35_manifest_free(manifest); + return 0; + } +} diff --git a/tests/test_qwen35_mlp_cpu_smoke.m b/tests/test_qwen35_mlp_cpu_smoke.m new file mode 100644 index 0000000..75487f2 --- /dev/null +++ b/tests/test_qwen35_mlp_cpu_smoke.m @@ -0,0 +1,116 @@ +#import +#import +#import +#import "../model/weight_loader.h" +#import "../kernels/inference/qwen_cpu_ops.h" + +static double abs_sum(const float *x, int n) { + double total = 0.0; + for (int i = 0; i < n; i++) { + total += fabs((double)x[i]); + } + return total; +} + +int main(int argc, const char *argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const char *blob_dir = argv[1]; + OrionQwen35Manifest *manifest = orion_qwen35_manifest_load(blob_dir); + if (!manifest) { + fprintf(stderr, "FAIL: manifest loader returned NULL\n"); + return 1; + } + + int d_model = manifest->d_model; + int d_ff = manifest->d_ff; + float *embed = (float *)calloc((size_t)d_model, sizeof(float)); + float *rms_weight = (float *)calloc((size_t)d_model, sizeof(float)); + float *normed = (float *)calloc((size_t)d_model, sizeof(float)); + float *mlp_out = (float *)calloc((size_t)d_model, sizeof(float)); + if (!embed || !rms_weight || !normed || !mlp_out) { + fprintf(stderr, "FAIL: allocation failed\n"); + goto fail; + } + + char path[2048]; + snprintf(path, sizeof(path), "%s/model/embed_tokens.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, d_model, embed)) { + fprintf(stderr, "FAIL: failed to read embedding row\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer0/input_layernorm.bin", blob_dir); + if (!orion_read_blob_row_f32(path, 0, d_model, rms_weight)) { + fprintf(stderr, "FAIL: failed to read layer0 RMSNorm\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer0/mlp_gate_proj.bin", blob_dir); + float *gate_proj = orion_read_blob_f32_exact(path, d_ff * d_model); + if (!gate_proj) { + fprintf(stderr, "FAIL: failed to read gate_proj\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer0/mlp_up_proj.bin", blob_dir); + float *up_proj = orion_read_blob_f32_exact(path, d_ff * d_model); + if (!up_proj) { + free(gate_proj); + fprintf(stderr, "FAIL: failed to read up_proj\n"); + goto fail; + } + + snprintf(path, sizeof(path), "%s/layer0/mlp_down_proj.bin", blob_dir); + float *down_proj = orion_read_blob_f32_exact(path, d_model * d_ff); + if (!down_proj) { + free(gate_proj); + free(up_proj); + fprintf(stderr, "FAIL: failed to read down_proj\n"); + goto fail; + } + + orion_qwen_cpu_rmsnorm(embed, rms_weight, d_model, 1e-6f, normed); + orion_qwen_cpu_swiglu_ffn(normed, gate_proj, up_proj, down_proj, d_model, d_ff, mlp_out); + + double embed_abs = abs_sum(embed, d_model); + double normed_abs = abs_sum(normed, d_model); + double mlp_out_abs = abs_sum(mlp_out, d_model); + + free(gate_proj); + free(up_proj); + free(down_proj); + + if (embed_abs <= 0.0 || normed_abs <= 0.0 || mlp_out_abs <= 0.0) { + fprintf(stderr, "FAIL: expected non-zero activation path\n"); + goto fail; + } + + printf("PASS: qwen35 mlp cpu smoke\n"); + printf(" blob_dir=%s\n", blob_dir); + printf(" token_id=%d\n", 0); + printf(" embed_abs_sum=%.6f\n", embed_abs); + printf(" normed_abs_sum=%.6f\n", normed_abs); + printf(" mlp_out_abs_sum=%.6f\n", mlp_out_abs); + printf(" next_blocker=%s\n", "implement attention kernels for linear/full layers"); + + free(embed); + free(rms_weight); + free(normed); + free(mlp_out); + orion_qwen35_manifest_free(manifest); + return 0; + +fail: + free(embed); + free(rms_weight); + free(normed); + free(mlp_out); + orion_qwen35_manifest_free(manifest); + return 1; + } +} diff --git a/tests/test_qwen35_prefill_frontend.m b/tests/test_qwen35_prefill_frontend.m new file mode 100644 index 0000000..a9668ac --- /dev/null +++ b/tests/test_qwen35_prefill_frontend.m @@ -0,0 +1,99 @@ +// tests/test_qwen35_prefill_frontend.m — Qwen3.5 ANE frontend smoke tests + +#import +#include "compiler/frontends/qwen35_prefill.h" +#include "compiler/model_config.h" +#include "compiler/pass_uniform_outputs.h" +#include "compiler/pipeline.h" +#include "compiler/validate.h" +#import "compiler/codegen.h" + +static int g_passed = 0, g_failed = 0; + +#define TEST(name) do { \ + printf(" %-40s ", #name); \ + if (test_##name()) { printf("PASS\n"); g_passed++; } \ + else { printf("FAIL\n"); g_failed++; } \ +} while(0) + +static const OrionModelConfig kQwen35 = { + .n_layer = 24, + .n_head = 8, + .n_kv_head = 2, + .d_model = 1024, + .head_dim = 256, + .hidden_dim = 3584, + .vocab = 248320, + .max_seq = 262144, +}; + +static NSString* compile_graph(OrionGraph* g) { + if (!g) return nil; + OrionValidationResult vr = orion_graph_validate(g); + if (!vr.valid) { + NSLog(@"graph validate failed: %s", vr.message); + orion_graph_free(g); + return nil; + } + orion_pipeline_optimize(g); + NSString* mil = orion_codegen_mil(g, "main"); + orion_graph_free(g); + return mil; +} + +static bool test_q_proj(void) { + OrionGraph* g = orion_frontend_qwen35_prefill_q_proj(3, 64, &kQwen35); + NSString* mil = compile_graph(g); + if (!mil) { + NSLog(@"q_proj compile_graph returned nil"); + return false; + } + return [mil containsString:@"self_attn_q_proj.bin"] && + [mil containsString:@"input_layernorm.bin"] && + [mil containsString:@"q_proj"]; +} + +static bool test_kv_proj(void) { + OrionGraph* g = orion_frontend_qwen35_prefill_kv_proj(3, 64, &kQwen35); + bool needs_uniform_pad = orion_pass_uniform_outputs(g); + if (needs_uniform_pad) { + NSLog(@"kv_proj unexpectedly needs uniform output padding"); + orion_graph_free(g); + return false; + } + NSString* mil = compile_graph(g); + if (!mil) { + NSLog(@"kv_proj compile_graph returned nil"); + return false; + } + return [mil containsString:@"self_attn_k_proj.bin"] && + [mil containsString:@"self_attn_v_proj.bin"] && + [mil containsString:@"k_proj"] && + [mil containsString:@"v_proj"]; +} + +static bool test_ffn(void) { + OrionGraph* g = orion_frontend_qwen35_prefill_ffn(3, 64, &kQwen35); + NSString* mil = compile_graph(g); + if (!mil) { + NSLog(@"ffn compile_graph returned nil"); + return false; + } + return [mil containsString:@"post_attention_layernorm.bin"] && + [mil containsString:@"mlp_gate_proj.bin"] && + [mil containsString:@"mlp_up_proj.bin"] && + [mil containsString:@"mlp_down_proj.bin"] && + [mil containsString:@"tanh("] && + [mil containsString:@"_sig"]; +} + +int main(int argc __attribute__((unused)), char* argv[] __attribute__((unused))) { + @autoreleasepool { + printf("test_qwen35_prefill_frontend:\n"); + TEST(q_proj); + TEST(kv_proj); + TEST(ffn); + printf("\n%d/%d passed\n", g_passed, g_passed + g_failed); + return g_failed > 0 ? 1 : 0; + } +} diff --git a/tests/test_qwen35_tokenizer_parity_smoke.m b/tests/test_qwen35_tokenizer_parity_smoke.m new file mode 100644 index 0000000..c011e59 --- /dev/null +++ b/tests/test_qwen35_tokenizer_parity_smoke.m @@ -0,0 +1,95 @@ +#import +#import "../tokenizer/gpt2_bpe.h" + +static NSDictionary* load_json(NSString* path) { + NSData* data = [NSData dataWithContentsOfFile:path]; + if (!data) return nil; + return [NSJSONSerialization JSONObjectWithData:data options:0 error:nil]; +} + +int main(int argc, const char* argv[]) { + @autoreleasepool { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + NSString* tokDir = [NSString stringWithUTF8String:argv[1]]; + NSString* vocabPath = [tokDir stringByAppendingPathComponent:@"vocab.json"]; + NSString* mergesPath = [tokDir stringByAppendingPathComponent:@"merges.txt"]; + NSString* metaPath = [tokDir stringByAppendingPathComponent:@"meta.json"]; + NSString* refsPath = [tokDir stringByAppendingPathComponent:@"parity_refs.json"]; + + NSDictionary* meta = load_json(metaPath); + NSDictionary* refs = load_json(refsPath); + if (!meta || !refs) { + fprintf(stderr, "FAIL: missing meta/parity refs\n"); + return 1; + } + + NSString* regex = meta[@"regex_pattern"]; + OrionGPT2Tokenizer* tok = orion_gpt2_tokenizer_load_with_regex( + vocabPath.UTF8String, + mergesPath.UTF8String, + regex.UTF8String + ); + if (!tok) { + fprintf(stderr, "FAIL: tokenizer load failed\n"); + return 1; + } + + NSArray* cases = refs[@"cases"]; + int encoded_ok = 0; + int decoded_ok = 0; + + for (NSDictionary* c in cases) { + NSString* text = c[@"text"]; + NSArray* ids = c[@"ids"]; + NSString* decoded_ref = c[@"decoded"]; + + int out[512] = {0}; + int n = orion_gpt2_encode(tok, text.UTF8String, out, 512); + if (n != (int)ids.count) { + fprintf(stderr, "FAIL: encode length mismatch for %s (%d vs %lu)\n", + text.UTF8String, n, (unsigned long)ids.count); + orion_gpt2_tokenizer_free(tok); + return 1; + } + for (int i = 0; i < n; i++) { + if (out[i] != [ids[i] intValue]) { + fprintf(stderr, "FAIL: token mismatch for %s at %d (%d vs %d)\n", + text.UTF8String, i, out[i], [ids[i] intValue]); + orion_gpt2_tokenizer_free(tok); + return 1; + } + } + encoded_ok++; + + char* decoded = orion_gpt2_decode(tok, out, n); + if (!decoded) { + fprintf(stderr, "FAIL: decode returned NULL\n"); + orion_gpt2_tokenizer_free(tok); + return 1; + } + NSString* decoded_ns = [NSString stringWithUTF8String:decoded]; + free(decoded); + if (![decoded_ns isEqualToString:decoded_ref]) { + fprintf(stderr, "FAIL: decode mismatch for %s\n", text.UTF8String); + fprintf(stderr, " got=%s\n", decoded_ns.UTF8String); + fprintf(stderr, " ref=%s\n", decoded_ref.UTF8String); + orion_gpt2_tokenizer_free(tok); + return 1; + } + decoded_ok++; + } + + printf("PASS: qwen35 tokenizer parity smoke\n"); + printf(" tokenizer_dir=%s\n", tokDir.UTF8String); + printf(" cases=%lu\n", (unsigned long)cases.count); + printf(" encoded_ok=%d\n", encoded_ok); + printf(" decoded_ok=%d\n", decoded_ok); + + orion_gpt2_tokenizer_free(tok); + return 0; + } +} diff --git a/tokenizer/gpt2_bpe.h b/tokenizer/gpt2_bpe.h index a578693..4848684 100644 --- a/tokenizer/gpt2_bpe.h +++ b/tokenizer/gpt2_bpe.h @@ -13,6 +13,13 @@ typedef struct OrionGPT2Tokenizer OrionGPT2Tokenizer; OrionGPT2Tokenizer* orion_gpt2_tokenizer_load(const char* vocab_path, const char* merges_path); +/// Load a ByteLevel BPE tokenizer with an explicit pre-tokenization regex. +/// This is used for non-GPT2 models such as Qwen that keep GPT2-style BPE +/// merges but alter the split regex. +OrionGPT2Tokenizer* orion_gpt2_tokenizer_load_with_regex(const char* vocab_path, + const char* merges_path, + const char* regex_pattern); + /// Encode text to token ids. /// @param tok Tokenizer handle /// @param text Input text (UTF-8) diff --git a/tokenizer/gpt2_bpe.m b/tokenizer/gpt2_bpe.m index 548ef78..6986296 100644 --- a/tokenizer/gpt2_bpe.m +++ b/tokenizer/gpt2_bpe.m @@ -62,6 +62,7 @@ static void build_unicode_to_byte(const unichar byte_to_unicode[256], uint8_t un unichar byte_to_unicode[256]; uint8_t unicode_to_byte[512]; void* bpe_cache; // NSMutableDictionary* + void* pretokenize_pattern; // NSString* }; // Accessors for type safety @@ -69,6 +70,7 @@ static void build_unicode_to_byte(const unichar byte_to_unicode[256], uint8_t un #define TOK_DECODER(tok) ((__bridge NSMutableArray*)(tok)->decoder) #define TOK_BPE_RANKS(tok) ((__bridge NSMutableDictionary*)(tok)->bpe_ranks) #define TOK_BPE_CACHE(tok) ((__bridge NSMutableDictionary*)(tok)->bpe_cache) +#define TOK_PRETOKENIZE_PATTERN(tok) ((__bridge NSString*)(tok)->pretokenize_pattern) #pragma mark - JSON Parsing (vocab.json) @@ -88,10 +90,15 @@ static bool load_vocab(OrionGPT2Tokenizer* tok, const char* path) { } NSMutableDictionary* enc = [NSMutableDictionary dictionaryWithCapacity:vocab.count]; - NSMutableArray* dec = [NSMutableArray arrayWithCapacity:vocab.count]; + int max_id = -1; + for (NSString* key in vocab) { + NSNumber* val = vocab[key]; + if (val.intValue > max_id) max_id = val.intValue; + } + NSMutableArray* dec = [NSMutableArray arrayWithCapacity:max_id + 1]; // Pre-fill decoder with empty strings - for (int i = 0; i < (int)vocab.count; i++) { + for (int i = 0; i <= max_id; i++) { [dec addObject:@""]; } @@ -106,7 +113,7 @@ static bool load_vocab(OrionGPT2Tokenizer* tok, const char* path) { tok->encoder = (void*)CFBridgingRetain(enc); tok->decoder = (void*)CFBridgingRetain(dec); - tok->vocab_size = (int)vocab.count; + tok->vocab_size = max_id + 1; return true; } @@ -148,19 +155,11 @@ static bool load_merges(OrionGPT2Tokenizer* tok, const char* path) { // GPT-2 pre-tokenization pattern: // 's|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+ // We implement this with NSRegularExpression. -static NSArray* pre_tokenize(NSString* text) { - static NSRegularExpression* regex = nil; - static dispatch_once_t onceToken; - dispatch_once(&onceToken, ^{ - NSString* pattern = @"'s|'t|'re|'ve|'m|'ll|'d" - @"| ?\\p{L}+" - @"| ?\\p{N}+" - @"| ?[^\\s\\p{L}\\p{N}]+" - @"|\\s+(?!\\S)" - @"|\\s+"; - regex = [NSRegularExpression regularExpressionWithPattern:pattern - options:0 error:nil]; - }); +static NSArray* pre_tokenize(OrionGPT2Tokenizer* tok, NSString* text) { + NSRegularExpression* regex = + [NSRegularExpression regularExpressionWithPattern:TOK_PRETOKENIZE_PATTERN(tok) + options:0 error:nil]; + if (!regex) return @[]; NSMutableArray* tokens = [NSMutableArray array]; NSArray* matches = @@ -263,8 +262,9 @@ static bool load_merges(OrionGPT2Tokenizer* tok, const char* path) { #pragma mark - Public API -OrionGPT2Tokenizer* orion_gpt2_tokenizer_load(const char* vocab_path, - const char* merges_path) { +OrionGPT2Tokenizer* orion_gpt2_tokenizer_load_with_regex(const char* vocab_path, + const char* merges_path, + const char* regex_pattern) { OrionGPT2Tokenizer* tok = calloc(1, sizeof(OrionGPT2Tokenizer)); build_byte_to_unicode(tok->byte_to_unicode); @@ -281,17 +281,39 @@ static bool load_merges(OrionGPT2Tokenizer* tok, const char* path) { } tok->bpe_cache = (void*)CFBridgingRetain([NSMutableDictionary dictionaryWithCapacity:10000]); + NSString* pattern = regex_pattern + ? [NSString stringWithUTF8String:regex_pattern] + : @"'s|'t|'re|'ve|'m|'ll|'d" + @"| ?\\p{L}+" + @"| ?\\p{N}+" + @"| ?[^\\s\\p{L}\\p{N}]+" + @"|\\s+(?!\\S)" + @"|\\s+"; + tok->pretokenize_pattern = (void*)CFBridgingRetain(pattern); fprintf(stderr, "tokenizer: loaded %d vocab, %d merges\n", tok->vocab_size, tok->num_merges); return tok; } +OrionGPT2Tokenizer* orion_gpt2_tokenizer_load(const char* vocab_path, + const char* merges_path) { + return orion_gpt2_tokenizer_load_with_regex(vocab_path, merges_path, NULL); +} + int orion_gpt2_encode(OrionGPT2Tokenizer* tok, const char* text, int* tokens, int max_tokens) { { NSString* nsText = [NSString stringWithUTF8String:text]; - NSArray* pre_tokens = pre_tokenize(nsText); + + // Honor exact added-token/special-token matches before pre-tokenization. + NSNumber* direct_id = TOK_ENCODER(tok)[nsText]; + if (direct_id && max_tokens > 0) { + tokens[0] = direct_id.intValue; + return 1; + } + + NSArray* pre_tokens = pre_tokenize(tok, nsText); int count = 0; for (NSString* chunk in pre_tokens) { @@ -349,5 +371,6 @@ void orion_gpt2_tokenizer_free(OrionGPT2Tokenizer* tok) { if (tok->decoder) CFRelease(tok->decoder); if (tok->bpe_ranks) CFRelease(tok->bpe_ranks); if (tok->bpe_cache) CFRelease(tok->bpe_cache); + if (tok->pretokenize_pattern) CFRelease(tok->pretokenize_pattern); free(tok); }