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
534 B
C
20 lines
534 B
C
// stories110m.h — Stories110M (Llama2-style, 12 layers, MHA)
|
|
#pragma once
|
|
|
|
#define MODEL_NAME "Stories110M"
|
|
|
|
#define DIM 768
|
|
#define HIDDEN 2048
|
|
#define HEADS 12
|
|
#define KV_HEADS 12
|
|
#define HD (DIM/HEADS) // = 64
|
|
#define GQA_RATIO 1 // MHA: no GQA
|
|
#define Q_DIM (HEADS * HD) // = 768 = DIM
|
|
#define KV_DIM (KV_HEADS * HD) // = 768 = DIM
|
|
#define SEQ 256
|
|
#define NLAYERS 12
|
|
#define VOCAB 32000
|
|
|
|
#define CKPT_PATH "ane_stories110M_dyn_ckpt.bin"
|
|
#define DEFAULT_DATA_PATH "../tinystories_data00.bin"
|