Exports one Spark decode step with the KV cache as explicit tensors, sized
per layer type: the 21 sliding layers hold a fixed 512 entries at any
context, the 7 full-attention layers grow. That asymmetry is the whole
reason the architecture is interesting on a phone, so the two types get
separately sized caches, masks and rotary angles.
Verified against the HF reference at 1024 context: 138.5 dB logits SNR,
identical top-1 and top-5.
Profiled on a real Galaxy S25 (Snapdragon 8 Elite) through AI Hub, QNN
context binary on the NPU, ms per layer:
sliding (512 KV, any ctx) fp16 1.689 w8a16 0.912 w4a16 0.917
full @ 1k fp16 1.559 w8a16 1.129 w4a16 1.127
full @ 4k fp16 2.964 w8a16 1.681 w4a16 1.662
Composed into a step (21 sliding + 7 full), w4a16 gives 27.2 ms at 1k and
30.9 ms at 4k, against 31.6 and 46.5 ms for the same model with every layer
full-attention -- 1.16x and 1.51x. Worth noting the hybrid is 0.94x at 1k in
fp16: sliding layers rotate all 256 head dims where full layers rotate 64,
and below roughly 2k that RoPE cost outweighs the shorter cache.
Also moves the head-wise attention gate off vLLM's CPU linear dispatch. The
gate is [2048, 8], and any linear that narrow is routed to the sgl-kernel
packed GEMM, whose fake-tensor rule assumes the packed weight keeps N in dim
0; VNNI packing does not at that size, so it reports K and torch.compile
fails tracing the gate multiply. A plain parameter with F.linear gains
nothing less -- a 2048x8 matmul has no use for packed AMX tiles -- and it
keeps the model compilable on CPU and GPU alike.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
