mirror of
https://github.com/hanzoai/ANE.git
synced 2026-10-30 18:53:51 +00:00
maderix
3f73962dc7
Implement Grouped-Query Attention (16q/8kv heads, head_dim=128) for Qwen3-0.6B (28 layers, 596M params). Model configs moved to models/*.h headers selected at build time via make MODEL=xxx. Key changes: - GQA-aware MIL kernels: sdpaFwd split from woFwd (Q_DIM!=DIM), qBwd/kvBwd split from qkvBwd (different IC dimensions) - K/V tile (KV_HEADS→HEADS) before SDPA backward, reduce after - 10 kernels total, all model-agnostic via compile-time defines - Makefile: make MODEL=qwen3_06b (default) or MODEL=stories110m - Both models verified: Stories110M ~115ms/step, Qwen3 ~412ms/step
20 lines
596 B
C
20 lines
596 B
C
// qwen3_06b.h — Qwen3-0.6B (28 layers, GQA 16q/8kv, head_dim=128)
|
|
#pragma once
|
|
|
|
#define MODEL_NAME "Qwen3-0.6B"
|
|
|
|
#define DIM 1024
|
|
#define HIDDEN 3072
|
|
#define HEADS 16
|
|
#define KV_HEADS 8
|
|
#define HD 128 // explicit head_dim (NOT DIM/HEADS)
|
|
#define GQA_RATIO (HEADS / KV_HEADS) // = 2
|
|
#define Q_DIM (HEADS * HD) // = 2048
|
|
#define KV_DIM (KV_HEADS * HD) // = 1024 (= DIM for this model)
|
|
#define SEQ 256
|
|
#define NLAYERS 28
|
|
#define VOCAB 151936
|
|
|
|
#define CKPT_PATH "ane_qwen3_06b_dyn_ckpt.bin"
|
|
#define DEFAULT_DATA_PATH "../tinystories_data00.bin"
|