Qwen3 written in UOps

qwen3.uir

# UNLIMITED
# Qwen3-0.6B in uops: greedy decoding of "The Mona Lisa was painted by" -> " Leonardo da Vinci."
# Prepare: python3 examples/prepare_qwen_uir.py   Run: python3 tests/test.py -k qwen

# --- weights: transposed to [in out] and stacked over the 28 layers ---

%embed     = buffer [1024 151936]     dtype=f32 slot=0 url=data/qwen/embed.bin
%qkv       = buffer [28 1024 32 2 64] dtype=f32 slot=1 url=data/qwen/qkv.bin
%qk_norm   = buffer [28 24 2 64]      dtype=f32 slot=2 url=data/qwen/qk_norm.bin
%o         = buffer [28 8 2 128 1024] dtype=f32 slot=3 url=data/qwen/o.bin
%gate_up   = buffer [28 1024 2 3072]  dtype=f32 slot=4 url=data/qwen/gate_up.bin
%down      = buffer [28 3072 1024]    dtype=f32 slot=5 url=data/qwen/down.bin
%attn_norm = buffer [28 1024]         dtype=f32 slot=6 url=data/qwen/attn_norm.bin
%mlp_norm  = buffer [28 1024]         dtype=f32 slot=7 url=data/qwen/mlp_norm.bin
%norm      = buffer 1024              dtype=f32 slot=8 url=data/qwen/norm.bin

# --- scalar functions ---

# e^x = (e^(x/256))^256, with a polynomial fit of e^r for |r| <= 0.35; below x = -87, e^x is ~0
%exp = {
  %r      = mul (max (param [] dtype=f32 slot=0 addrspace=ALU) -87.0) 0.00390625
  %p      = add (mul %r (add (mul %r (add (mul %r 0.00836916) 0.04187572)) 0.16666505)) 0.49999371
  %e      = add (mul %r (add (mul %r %p) 1.0)) 1.0
  %e4     = mul (mul %e %e) (mul %e %e)
  %e16    = mul (mul %e4 %e4) (mul %e4 %e4)
  %e64    = mul (mul %e16 %e16) (mul %e16 %e16)
  %result = mul (mul %e64 %e64) (mul %e64 %e64)
}

# sin(x) for x >= 0: remove whole turns so y is in [-pi, pi], then a polynomial fit of sin(y)
%sine = {
  %x      = param [] dtype=f32 slot=0 addrspace=ALU
  %y      = add %x (mul (trunc (add (mul %x 0.15915494) 0.5)) -6.2831855)
  %y2     = mul %y %y
  %p      = add (mul %y2 (add (mul %y2 -2.1381794e-8) 2.7248525e-6)) -0.00019830526
  %result = add %y (mul (mul %y %y2) (add (mul %y2 (add (mul %y2 %p) 0.0083332015)) -0.16666664))
}

# 1/sqrt(sum_sq/n + 1e-6), the rmsnorm scale, by Heron's method: s = (s + m/s)/2 from s = m
%inv_rms = {
  %sum_sq = param [] dtype=f32 slot=0 addrspace=ALU
  %inv_n  = param [] dtype=f32 slot=1 addrspace=ALU
  %m      = add (mul %sum_sq %inv_n) 1e-6
  %s      = alloc [] dtype=f32 slot=0
  %start  = store %s %m
  %k      = range 24 slot=0
  %prev   = load (after %s %start %k)
  %steps  = end (store %s (mul (add %prev (mul %m (reciprocal %prev))) 0.5)) %k
  %result = reciprocal (load (after %s %steps))
}

# --- decoding: each position embeds its token, runs the layers and predicts the next token ---

%tokens   = buffer 10 dtype=i32 slot=9
%prompt   = after %tokens (store %tokens [785 98783 28556 572 23983 553 0 0 0 0])
%keys     = alloc [28 9 8 2 64] dtype=f32 slot=0
%values   = alloc [28 9 8 2 64] dtype=f32 slot=1
%x        = alloc 1024 dtype=f32 slot=2
%pos      = range 9 slot=0
%token    = floormod (load (index (after %prompt %pos) %pos)) 151936
%embedded = store %x (load (index (permute %embed [1 0]) %token))

# rotary embedding: cos and sin of pos * theta^(-2i/128), with theta = 1e6 and i < 64
%freq  = range 64 slot=1
%angle = mul (cast %pos dtype=f32) (call %exp (mul (cast %freq dtype=f32) -0.21586735))
%rot   = permute (load (stage [(call %sine (add %angle 1.5707964)) (call %sine %angle)] %freq)) [1 0]
%cos   = index %rot 0
%sin   = index %rot 1

# --- one layer: x += attention(rmsnorm(x)), then x += mlp(rmsnorm(x)) ---

%layer  = range 28 slot=2
%h      = load (after %x %embedded %layer)
%rms1   = call %inv_rms (reduce (mul %h %h) op=add pop=1) 0.0009765625
%n1     = mul (mul %h %rms1) (load (index %attn_norm %layer))
%heads  = reduce (mul (reshape %n1 [1024 1 1 1]) (load (index %qkv %layer))) op=add pop=1

# q and k (16 + 8 heads, each two halves of 64): rmsnorm per head, then rotate the halves.
# Using n = sqrt(128) instead of 128 scales q and k by 128^-1/4 each: the attention's 1/sqrt(128).
%qk     = shrink %heads [0 0 0] [24 2 64]
%head   = range 24 slot=3
%sumsq  = reduce (permute (mul %qk %qk) [1 2 0]) op=add pop=2
%scale  = load (stage [(call %inv_rms (index %sumsq %head) 0.088388346)] %head)
%halves = mul (permute (mul %qk (load (index %qk_norm %layer))) [1 0 2]) %scale
%x1     = index %halves 0
%x2     = index %halves 1
%rot1   = add (mul %x1 %cos) (mul (mul %x2 %sin) -1.0)
%rot2   = add (mul %x2 %cos) (mul %x1 %sin)
%qk_r   = permute [%rot1 %rot2] [1 0 2]

# attention over the cached positions t <= pos, where query head m reads k/v head m/2.
# Scores stay far below 88, so softmax needs no max subtraction; masked scores are -87 (e^-87 ~ 0).
%store_k = store (index %keys %layer %pos) (shrink %qk_r [16 0 0] [8 2 64])
%store_v = store (index %values %layer %pos) (shrink %heads [24 0 0] [8 2 64])
%k_all   = load (index (after %keys %store_k) %layer)
%v_all   = reshape (load (index (after %values %store_v) %layer)) [9 8 1 128 1]
%t       = range 9 slot=4
%m       = range 16 slot=5
%score   = reduce (mul (index %qk_r %m) (index %k_all %t (floordiv %m 2))) op=add pop=2
%weight  = load (stage (call %exp (where (cmplt %pos %t) -87.0 %score)) %t %m)
%probs   = reshape (mul %weight (reciprocal (reduce %weight op=add pop=1))) [9 8 2 1 1]
%attn    = reduce (mul %probs %v_all) op=add pop=1
%h1      = add %h (reduce (mul %attn (load (index %o %layer))) op=add pop=3)

# mlp: down(silu(gate) * up), with silu(g) = g / (1 + e^-g)
%rms2   = call %inv_rms (reduce (mul %h1 %h1) op=add pop=1) 0.0009765625
%n2     = mul (mul %h1 %rms2) (load (index %mlp_norm %layer))
%gu     = reduce (mul (reshape %n2 [1024 1 1]) (load (index %gate_up %layer))) op=add pop=1
%f      = range 3072 slot=6
%gate   = index %gu 0 %f
%silu   = mul %gate (reciprocal (add (call %exp (mul %gate -1.0)) 1.0))
%act    = load (stage [(mul %silu (index %gu 1 %f))] %f)
%h2     = add %h1 (reduce (mul %act (load (index %down %layer))) op=add pop=1)
%layers = end (store %x %h2) %layer

# --- next token: the argmax of embed @ rmsnorm(x), ties picking the highest id ---
# predictions at prompt positions land in slot 6, until position 5 writes the first real one there

%hf     = load (after %x %layers)
%rmsf   = call %inv_rms (reduce (mul %hf %hf) op=add pop=1) 0.0009765625
%nf     = mul (mul %hf %rmsf) (load %norm)
%logits = reduce (mul (reshape %nf [1024 1]) (load %embed)) op=add pop=1
%id     = range 151936 slot=7
%next   = reduce (where (cmplt (index %logits %id) (reduce %logits op=max pop=1)) 0 %id) %id op=max
%write  = store (index (after %prompt %pos) (max (add %pos 1) 6)) %next
%0      = assert (after %tokens (end %write %pos)) [785 98783 28556 572 23983 553 65386 2994 96766 13]