qwen3: 0.6B correctness + 2.6x decode speedup - #7
Merged
Merged
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e76138fcb4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| float rope_beta_slow; | ||
|
|
||
| float eps; | ||
| uint32_t tie_word_embeddings; // 1 = tie LM head to embedding (Qwen3-0.6B), 0 = separate lm_head (Qwen3-8B) |
There was a problem hiding this comment.
Adding tie_word_embeddings to sk_qwen_config changes the C struct layout, but models/qwen/qwen.py still defines _Config with fields only through eps. When Python passes this shorter struct to sk_qwen_create, h->cfg = *cfg reads past the ctypes buffer, so tie_word_embeddings becomes undefined and model behavior (tied vs untied LM head path) becomes nondeterministic for all Python callers.
Useful? React with 👍 / 👎.
| const void* w_gate; | ||
| const void* w_up; | ||
| const void* w_down; | ||
| const void* w_lm_head; // optional (untied); may be null when tie_word_embeddings=1 |
There was a problem hiding this comment.
Adding w_lm_head to sk_qwen_weights also breaks the existing ctypes binding: models/qwen/qwen.py builds _Weights from _WEIGHT_FIELDS that do not include this member. sk_qwen_load_weights now reads w->w_lm_head from memory beyond the Python struct, and if that garbage value is non-null the memcpy path can dereference an invalid address and crash.
Useful? React with 👍 / 👎.
| if (!hp || !out_fp16) return -1; | ||
| auto* h = reinterpret_cast<meow::qwen::Handle*>(hp); | ||
| std::memcpy(out_fp16, h->bufs.capture->contents(), h->bufs.capture->length()); | ||
| return 0; |
There was a problem hiding this comment.
sk_qwen_get_capture always returns success and copies the capture buffer, but the API comment says it should return -1 when the requested layer was not captured. As implemented, calling get_capture before any matching forward pass (or with an out-of-range layer index) returns stale/zero data as if valid, which can silently corrupt debugging or evaluation workflows.
Useful? React with 👍 / 👎.
…memory determinism bug The Python _Config ctypes struct was missing the tie_word_embeddings field added in e76138f, so the C side read uninitialized memory. Garbage nonzero values flipped Qwen3-0.6B into untied-lm_head mode, allocating an unloaded zero buffer for w_lm_head and producing nondeterministic post-loop logits / silent hangs in layer_diff at L>=1. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…, not whole T_max buffer) Previously memcpy'd the full capture buffer (T_max * d_model fp16) into the caller's output slot, smashing the stack/heap whenever the caller sized for the actual seq used at forward (e.g. T=2 → 4KB) instead of seq_max (e.g. 128 → 256KB). This was causing SIGSEGV in layer_diff_06b after the first cap_layer() call. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two compounding bugs that made SK Qwen3 diverge from HF reference: 1. `gated_mlp` Metal kernel was structurally broken: its threadgroup `intermediate` scratch had shape (BM, BN=64) but the down-projection reduction loop iterated `bk` over all of N_int (e.g. 3072), reading the same 32 stale intermediate columns N_int/BK times. Also the write index used absolute `gc` instead of `gc % BN`, OOB-writing in every TG with col>0. Replaced its invocation in the qwen launcher with three separate `gemm_fp16` passes (gate, up, down) plus a tiny new `silu_mul_f16` elementwise kernel. 2. `qwen_rope_qk` writes Q/K in seq-major (T, H, D) layout, but `kv_cache_write` and `mha_causal` expect head-major (H, T, D). With T=1 the two layouts coincide so single-token decode appeared to work; prefill with T>=2 scrambled K-cache rows across heads/positions and blew up to fp16 saturation by layer 2. Added `transpose_seq_to_head_f16` / `transpose_head_to_seq_f16` kernels and four transpose calls per layer: Q/K/V seq→head after RoPE, attn_out head→seq before o_proj. Per-layer residual rel_err (Qwen3-0.6B, "Hi!" prompt, T=2): before: L0=0.51, L1=0.76, L2..L27 ~ 1.0 (fp16 saturation) after: L0=0.0018, L1=0.0031, L27=0.0172 SK last-token argmax = 358 (matches HF). Top-5 identical to HF. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Previously dispatch_layer encoded every GEMM (QKV, O, gate, up, down) through the tile-MMA gemm_fp16 kernel even at decode (M=1), wasting 31/32 rows of every BM=32 tile. Added a gemv_m1 PSO to LayerPSOs and taught encode_gemm to switch to the specialized M=1 matvec when the caller passes the optional pso. Same trick the gemm.c++ host wrapper already uses for ad-hoc calls; it just wasn't wired into the model launcher. bench (Qwen3-0.6B decode, lexie M4): before: 27.3 tok/s after: 64.0 tok/s (+134%) correctness preserved (full 28-layer diff < 2% rel, argmax=358, top-5 identical to HF). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
At seq==1 the (T,H,D) and (H,T,D) layouts are bytewise identical, so the 4 transpose copies per layer are pure waste. Alias the source buffers into kv_cache_write / mha_causal / o_proj instead. Prefill (T>=2) path unchanged. bench: 64.0 → 64.4 tok/s (Qwen3-0.6B decode on lexie M4) correctness: argmax=358, top-5 identical to HF (decode path). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two M=1 decode-path kernels: 1. gemv_swiglu_fp16_m1: fuses gate-GEMV, up-GEMV, and silu_mul into one dispatch. m_in is loaded once into TG memory and used to dot against both W_gate and W_up per column; SiLU(gate)*up is written to the output buffer directly. Replaces 3 dispatches with 1 in the MLP critical path (no speedup observed in isolation — we are bandwidth-bound on the weight reads — but reduces encode overhead and frees an intermediate buffer round-trip). 2. gemv_t_fp16_m1: M=1 matvec for transB=1 (LM-head case where weights are stored (V, D) row-major). Each thread owns one output row and walks its weight vector contiguously. Replaces the LM-head GEMM (151936 × 1024 fp16, 311 MB) which had been falling back to the tile-MMA gemm_fp16 path at M=1. bench (Qwen3-0.6B decode, lexie M4): prior: 64.4 tok/s now: 71.0 tok/s (+10%) correctness preserved (argmax=358, top-5 identical to HF). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Add Dtype::Q8_0 to weight_store. - Map ggml type 8 → Dtype::Q8_0 in gguf loader. - New kernel kernels/gemm/q8_0_matvec.metal: simdgroup-reduced Q8_0 weight × fp16 activation matvec for decode (M=1). NSG=4 SGs/TG, NR0=2 rows/SG, NQ=8 acts/lane; tiles K dimension in blocks of NSG*NQ*32. Output fp16. Launcher/Qwen dispatch wiring NOT yet plumbed; this commit only lands the prerequisites so the kernel binary and dtype tag are available.
Changes:
- inference/weight_store.h: add dtype_bytes(d, n_elems) helper.
- kernels/gemm/q8_0_matvec.metal: fix dispatch contract — each TG writes
NR0=2 rows (not NSG*NR0=8). All 4 simdgroups cooperate over K for the
same 2 rows. The previous comment claimed 8 rows/TG which caused the
host dispatch to undercount TGs and leave 3/4 of the output zero.
Also guard final simd_sum so all 32 lanes participate.
- models/qwen/qwen_model.h: encode_gemm_qaware routes M=1 (decode) and
M>1 (prefill, per-row loop) through q8_0_matvec when dt_w==Q8_0. Per-
projection dtype fields on LayerBuffers/ModelWeights; per-layer byte
offsets use dtype_bytes. LM-head also routes to q8_0_matvec when
dt_lm_head==Q8_0. Skip fused SwiGLU GEMV when MLP weights are Q8_0.
- models/qwen/launcher.c++: register q8_0_matvec PSO (nullptr-tolerant).
- models/qwen/weights.{c++,h}: sk_qwen_load_gguf parses GGUF, dequants
norms/embed to fp16, reallocates 6 projection buffers to Q8_0 size and
memcpys raw blocks, allocates Q8_0 lm_head from token_embd (tied) or
output.weight (untied). Exposes sk_qwen_debug_q8_matvec for unit tests.
- models/qwen/qwen.py: Qwen.load_gguf(path) helper.
Validation (Qwen3-0.6B, M1):
- last-token argmax ("Hi!"): Q8_0 = 358 = FP16 = HF.
- L0 residual rel_err 1.06%, logits rel_err 4.51%.
- Decode: 136.1 tok/s (fp16 baseline 71.7 -> 1.90x; >=120 target met;
llama.cpp Q8_0 reference 118.8).
- Standalone q8_0_matvec unit test: rel_err 0.59% vs numpy reference.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Qwen now inherits SuperKittens.inference.generation.Model and implements the _forward / _last_logits / reset contract so the shared generate() / chat() loop works. - Add bake_and_set_rope() and bind sk_qwen_set_rope_tables / sk_qwen_get_last_logits in the ctypes shim. - New SuperKittens/models/qwen/__init__.py registers "qwen3-0.6b" with sk.load, derives Config from the HF snapshot config.json, loads Q8_0 GGUF (or fp16 safetensors), bakes RoPE tables, and attaches a tokenizer via the HF tokenizer.json with the qwen3 chat template. - Qwen.chat() accepts a string or messages list and applies the chat template by default. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Rename the factory parameter spec -> variant so it doesn't clash with the positional spec argument used by sk.api.register(). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Qwen3 0.6B SK port — correct, +2.6x decode.
next: Q8_0 quant to close the last 40% (agent in flight).