From 62a3ab394bb04d5c8f926122f0cc3360229b7b7c Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Tue, 9 Jun 2026 21:19:08 -0400 Subject: [PATCH 01/31] =?UTF-8?q?preserve:=20gemma4=20batched-lane=20prefi?= =?UTF-8?q?ll=20WIP=20(agent=20context-death=20recovery;=20lane-divergence?= =?UTF-8?q?=20unresolved=20=E2=80=94=20no=20cross-lane=20bleed=20per=20per?= =?UTF-8?q?m=20test,=20A/B=20both=20degenerate=20on=20diverging=20real-pro?= =?UTF-8?q?mpt=20lanes)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SuperKittens/models/gemma/gemma4/gemma4.py | 26 +++++ .../models/gemma/gemma4/gemma4_model.h | 98 +++++++++++++++++++ SuperKittens/models/gemma/gemma4/launcher.c++ | 94 ++++++++++++++++++ SuperKittens/models/gemma/gemma4/launcher.h | 11 +++ 4 files changed, 229 insertions(+) diff --git a/SuperKittens/models/gemma/gemma4/gemma4.py b/SuperKittens/models/gemma/gemma4/gemma4.py index d9f6ba3..5c7bae3 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4.py +++ b/SuperKittens/models/gemma/gemma4/gemma4.py @@ -242,6 +242,32 @@ def forward_batched(self, input_ids: np.ndarray) -> np.ndarray: raise RuntimeError(f"sk_gemma4_forward_batched failed: {ret}") return out + def prefill_batched(self, input_ids: np.ndarray, chunk_size: int = 0) -> np.ndarray: + """Batched chunked prefill: input_ids is (batch, seq) int32 (request-major, + equal-length lockstep lanes); returns (batch,) int32 greedy next tokens. + One M=batch*chunk GEMM pass per chunk amortizes the weight stream across + lanes AND tokens (vs. forward_batched seq=1 x T). chunk_size 0 -> seq_max. + Token-exact while seq <= cfg.window (local-attention ring eviction).""" + ids = np.ascontiguousarray(input_ids, dtype=np.int32) + if ids.ndim != 2 or ids.shape[0] != self.cfg.batch: + raise ValueError(f"need (batch={self.cfg.batch}, seq) ids, got {ids.shape}") + fn = getattr(_load(), "sk_gemma4_prefill_batched", None) + if fn is None: + raise RuntimeError("libsk.dylib has no sk_gemma4_prefill_batched; rebuild dylib") + fn.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.c_uint32, + ctypes.POINTER(ctypes.c_int32)] + fn.restype = ctypes.c_int + out = np.empty((self.cfg.batch,), dtype=np.int32) + ret = fn(self._h, + ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)), + ctypes.c_uint32(ids.shape[1]), + ctypes.c_uint32(chunk_size), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))) + if ret: + raise RuntimeError(f"sk_gemma4_prefill_batched failed: {ret}") + return out + def set_dump_enabled(self, enabled: bool): _load().sk_gemma4_set_dump_enabled(self._h, 1 if enabled else 0) diff --git a/SuperKittens/models/gemma/gemma4/gemma4_model.h b/SuperKittens/models/gemma/gemma4/gemma4_model.h index e834239..8fbc989 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4_model.h +++ b/SuperKittens/models/gemma/gemma4/gemma4_model.h @@ -1158,6 +1158,16 @@ struct ModelParams { // rows so each request gets its own next token. Default 0 keeps the // single-row (last-position) decode/prefill path byte-identical. uint32_t decode_all_rows = 0; + + // Batched (batch=N, seq>1) chunked prefill. 0 = off (all existing paths + // byte-identical). 1 = interior chunk: layers/KV only, skip final norm + + // head + descale + softcap + argmax (serving needs logits only after the + // full prompt). 2 = final chunk: project each lane's LAST prompt row + // (b*seq+seq-1) -> logits row b, then descale + softcap + argmax -> + // output_id[b]. The decode_all_rows head would project all batch*seq rows + // (vocab x T dead work) and its argmax requires seq==1, so the prefill + // tail is its own path. + uint32_t batched_prefill = 0; }; struct ModelPSOs { @@ -1755,6 +1765,9 @@ inline void dispatch_model( MTL::Buffer* tmp = cur; cur = nxt; nxt = tmp; } + // Interior batched-prefill chunk: KV is written, no logits consumer yet. + if (M.batched_prefill == 1u) return; + // C. Final RMSNorm { auto* enc = cmd->computeCommandEncoder(); @@ -1782,6 +1795,91 @@ inline void dispatch_model( _dump_blit_row(cmd, nxt, B.dump_stash, T, M.d_model, base * M.d_model); } + // Final batched-prefill chunk: per-lane last-row head + descale + softcap + + // argmax (see the batched_prefill field WHY). The lane matvecs have + // disjoint outputs and a read-only head weight, so one encoder serves all + // lanes; encoder boundaries order head -> descale -> softcap -> argmax. + if (M.batched_prefill == 2u) { + const uint32_t K_v = M.d_model; + const uint32_t N_v = M.vocab_size; + const size_t in_row_bytes = (size_t)M.d_model * 2; + const size_t out_row_bytes = (size_t)M.vocab_size * 2; + const bool q4k_head = (W.w_lm_head_q4k != nullptr) && (P.layer.q4k_matvec_bf16 != nullptr); + const bool q8_head = (W.w_lm_head_q8 != nullptr) && (P.layer.q8_0_matvec_bf16 != nullptr); + { + auto* enc = cmd->computeCommandEncoder(); + for (uint32_t b = 0; b < M.batch; ++b) { + const size_t off_A = ((size_t)b * M.seq + M.seq - 1u) * in_row_bytes; + const size_t off_C = (size_t)b * out_row_bytes; + if (q4k_head || q8_head) { + const uint32_t NR0 = 2; + enc->setComputePipelineState(q4k_head ? P.layer.q4k_matvec_bf16 + : P.layer.q8_0_matvec_bf16); + enc->setBuffer(nxt, off_A, 0); + enc->setBuffer(q4k_head ? W.w_lm_head_q4k : W.w_lm_head_q8, 0, 1); + enc->setBuffer(B.logits, off_C, 2); + enc->setBytes(&K_v, 4, 3); + enc->setBytes(&N_v, 4, 4); + enc->dispatchThreadgroups(MTL::Size((N_v + NR0 - 1) / NR0, 1, 1), + MTL::Size(128, 1, 1)); + } else { + const uint32_t M_v = 1u; + uint32_t ldA = K_v, ldB = K_v, ldC = N_v; + int transA = 0, transB = 1, has_bias = 0; + enc->setComputePipelineState(P.layer.gemm); + enc->setBuffer(nxt, off_A, 0); + enc->setBuffer(W.w_embed, 0, 1); + enc->setBuffer(B.logits, off_C, 2); + enc->setBytes(&M_v, 4, 3); enc->setBytes(&N_v, 4, 4); + enc->setBytes(&K_v, 4, 5); enc->setBytes(&ldA, 4, 6); + enc->setBytes(&ldB, 4, 7); enc->setBytes(&ldC, 4, 8); + enc->setBytes(&transA, 4, 9); enc->setBytes(&transB, 4, 10); + enc->setBytes(&has_bias, 4, 11); + enc->setBuffer(B.logits, off_C, 12); + enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, 1, 1), + MTL::Size(64, 1, 1)); + } + } + enc->endEncoding(); + } + // Lane logits live in rows 0..batch — descale/softcap span batch*vocab. + const uint32_t n_lane = M.batch * M.vocab_size; + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.logit_descale); + enc->setBuffer(B.logits, 0, 0); + float inv_scale = 1.0f / std::sqrt((float)M.d_model); + enc->setBytes(&n_lane, 4, 1); + enc->setBytes(&inv_scale, 4, 2); + uint32_t groups = ((n_lane / 4u) + 127u) / 128u; + enc->dispatchThreadgroups(MTL::Size(groups, 1, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + if (M.final_logit_softcap > 0.0f) { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.logit_softcap); + enc->setBuffer(B.logits, 0, 0); + float cap = M.final_logit_softcap; + enc->setBytes(&n_lane, 4, 1); + enc->setBytes(&cap, 4, 2); + uint32_t groups = ((n_lane / 4u) + 127u) / 128u; + enc->dispatchThreadgroups(MTL::Size(groups, 1, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.argmax); + for (uint32_t b = 0; b < M.batch; ++b) { + enc->setBuffer(B.logits, (size_t)b * out_row_bytes, 0); + enc->setBuffer(B.output_id, (size_t)b * sizeof(int32_t), 1); + enc->setBytes(&M.vocab_size, 4, 2); + enc->dispatchThreadgroups(MTL::Size(1, 1, 1), MTL::Size(1024, 1, 1)); + } + enc->endEncoding(); + } + return; + } + // D. LM head GEMM (tied with input embedding). // // Q8_0 fast path (decode, T=1): when w_lm_head_q8 is populated and the diff --git a/SuperKittens/models/gemma/gemma4/launcher.c++ b/SuperKittens/models/gemma/gemma4/launcher.c++ index 6e7676a..1c4b039 100644 --- a/SuperKittens/models/gemma/gemma4/launcher.c++ +++ b/SuperKittens/models/gemma/gemma4/launcher.c++ @@ -681,6 +681,100 @@ extern "C" int sk_gemma4_forward_batched(sk_gemma4_handle* hp, return 0; } +// Batched chunked prefill: ids is batch*seq int32 (request-major: row b is lane +// b's seq prompt tokens; all lanes the same length, lockstep positions). Runs +// the prompt in chunks of <= chunk_size (clamped to seq_max): each chunk's +// projections are ONE M=batch*chunk GEMM pass, so the weight stream is +// amortized across lanes AND tokens — vs. the token-by-token serving prefill +// that pays `seq` full weight-read passes. Logits (per-lane last row, with +// gemma's descale + softcap) and argmax run only on the final chunk; +// out_next[b] receives lane b's greedy next token. chunk_size 0 -> seq_max. +// Does NOT reset; lanes share current_pos (lockstep), KV per lane via its +// cache slice. SWA exactness bound (same as the single-stream seq>1 path): +// local-layer KV is a ring of size `window`, so a chunk's tail overwrites the +// oldest in-window keys of its own interior rows once current_pos+seq exceeds +// window — prompts must satisfy seq <= window for token-exact prefill. +extern "C" int sk_gemma4_prefill_batched(sk_gemma4_handle* hp, + const int* ids, uint32_t seq, + uint32_t chunk_size, int* out_next) { + if (!hp || !ids || !out_next) return -1; + auto* h = reinterpret_cast(hp); + if (seq == 0) return -2; + if (h->cfg.batch < 1) return -5; + uint32_t step = (chunk_size == 0) ? h->cfg.seq_max : chunk_size; + if (step > h->cfg.seq_max) step = h->cfg.seq_max; + if (h->current_pos + seq > h->cfg.cache_max) return -4; + + auto* dev = sk::bindings_device(); + auto* q = sk::bindings_queue(); + if (!dev || !q) return -3; + + int32_t* in_ids = (int32_t*)h->bufs.input_ids->contents(); + const uint32_t batch = h->cfg.batch; + for (uint32_t s = 0; s < seq; ) { + const uint32_t remaining = seq - s; + const uint32_t this_seq = (remaining < step) ? remaining : step; + for (uint32_t b = 0; b < batch; ++b) + std::memcpy(in_ids + (size_t)b * this_seq, ids + (size_t)b * seq + s, + (size_t)this_seq * sizeof(int32_t)); + const bool final_chunk = (s + this_seq == seq); + + meow::gemma4::ModelParams mp; + mp.batch = h->cfg.batch; + mp.seq = this_seq; + mp.n_layers = h->cfg.n_layers; + mp.local_period = h->cfg.local_period; + mp.d_model = h->cfg.d_model; + mp.n_int = h->cfg.n_int; + mp.n_heads = h->cfg.n_heads; + mp.n_kv_heads_local = h->cfg.n_kv_heads_local; + mp.n_kv_heads_global = h->cfg.n_kv_heads_global; + mp.head_dim_local = h->cfg.head_dim_local; + mp.head_dim_global = h->cfg.head_dim_global; + mp.window = h->cfg.window; + mp.cache_max = h->cfg.cache_max; + mp.prope_p_pairs = h->cfg.prope_p_pairs; + mp.vocab_size = h->cfg.vocab_size; + mp.ple_dim = h->cfg.ple_dim; + mp.has_ple = (h->cfg.has_ple != 0); + mp.full_rope_global = (h->cfg.full_rope_global != 0); + mp.apply_layer_scalar = (h->cfg.apply_layer_scalar != 0); + mp.eps = h->cfg.eps; + mp.final_logit_softcap = h->cfg.final_logit_softcap; + mp.current_pos = h->current_pos; + mp.n_int_per_layer = h->n_int_per_layer.data(); + mp.mlp_gate_off_e = h->mlp_gate_off_e.data(); + mp.mlp_down_off_e = h->mlp_down_off_e.data(); + mp.kv_source_layer = h->kv_source_layer.data(); + mp.layer_scalar_host = h->layer_scalar_host.data(); + mp.dump_enabled = h->dump_enabled; + mp.batched_prefill = final_chunk ? 2u : 1u; + + auto* cmd = q->commandBuffer(); + meow::gemma4::dispatch_model(cmd, h->psos, h->weights, h->bufs, mp); + cmd->commit(); + cmd->waitUntilCompleted(); + if (std::getenv("SK_GEMMA4_GPUPROF")) + std::fprintf(stderr, "[gpuprof] gpu_busy_us=%.1f\n", + (cmd->GPUEndTime() - cmd->GPUStartTime()) * 1e6); + if (cmd->status() == MTL::CommandBufferStatusError) { + auto* e = cmd->error(); + std::fprintf(stderr, "gemma4 prefill_batched: command buffer ERROR (status=%ld): %s\n", + (long)(e ? e->code() : -1), + e && e->localizedDescription() ? e->localizedDescription()->utf8String() : "?"); + cmd->release(); + return -7; + } + cmd->release(); + h->current_pos += this_seq; + s += this_seq; + } + + std::memcpy(out_next, h->bufs.output_id->contents(), + (size_t)batch * sizeof(int32_t)); + return 0; +} + extern "C" int sk_gemma4_get_last_logits(sk_gemma4_handle* hp, void* out_fp16) { if (!hp || !out_fp16) return -1; auto* h = reinterpret_cast(hp); diff --git a/SuperKittens/models/gemma/gemma4/launcher.h b/SuperKittens/models/gemma/gemma4/launcher.h index bfe0873..6465cf0 100644 --- a/SuperKittens/models/gemma/gemma4/launcher.h +++ b/SuperKittens/models/gemma/gemma4/launcher.h @@ -100,6 +100,17 @@ int sk_gemma4_forward_batched(sk_gemma4_handle* h, const int* input_ids, uint32_t seq, int* output_id); +// Batched (batch=N, seq>1) chunked prefill. ids is batch*seq int32 +// (request-major; all lanes the same length, lockstep positions). Runs chunks +// of <= chunk_size (0 -> seq_max; clamped to seq_max) with one M=batch*chunk +// GEMM pass per chunk; logits + argmax only on the final chunk (per-lane LAST +// row, with gemma's logit descale + softcap). out_next receives `batch` greedy +// next tokens. Advances current_pos by seq; does NOT reset. Token-exact while +// seq <= window (local-layer ring eviction; see launcher.c++ WHY). +int sk_gemma4_prefill_batched(sk_gemma4_handle* h, + const int* ids, uint32_t seq, + uint32_t chunk_size, int* out_next); + // Reset the per-sequence KV-cache cursor (does NOT zero the cache buffers; // stale data is shadowed once new K/V is written over it). void sk_gemma4_reset(sk_gemma4_handle* h); From bf05198e56f0e8eda9de08c5f5bbdbc027f54049 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 14:55:25 -0400 Subject: [PATCH 02/31] =?UTF-8?q?gemm:=20fix=20fp16-GEMM=20row=20grid=20at?= =?UTF-8?q?=20remaining=20/64=20call=20sites=20(qwen=20lm=5Fhead=20fallbac?= =?UTF-8?q?k=20+=20generic=20gemm=20API)=20=E2=80=94=20gemm=5Ffp16=20tiles?= =?UTF-8?q?=20BM=3D32;=20M>32=20skipped=20rows=2032..63=20per=20block?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SuperKittens/kernels/gemm/gemm.c++ | 2 +- SuperKittens/models/qwen/qwen_model.h | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/SuperKittens/kernels/gemm/gemm.c++ b/SuperKittens/kernels/gemm/gemm.c++ index 1bf48a6..4227216 100644 --- a/SuperKittens/kernels/gemm/gemm.c++ +++ b/SuperKittens/kernels/gemm/gemm.c++ @@ -16,7 +16,7 @@ static int dispatch(const char* kname, void* A, void* B, void* C, void* bias, auto* bBias = has_bias ? sk::bindings_device()->newBuffer(bbb, MTL::ResourceStorageModeShared) : nullptr; if (bBias) memcpy(bBias->contents(), bias, bbb); - uint32_t gx = (N + 63) / 64, gy = (M + 63) / 64; + uint32_t gx = (N + 63) / 64, gy = (M + 31) / 32; // gemm_fp16 BM=32 rows auto* cmd = sk::bindings_queue()->commandBuffer(); auto* enc = cmd->computeCommandEncoder(); diff --git a/SuperKittens/models/qwen/qwen_model.h b/SuperKittens/models/qwen/qwen_model.h index fc2f872..f92991c 100644 --- a/SuperKittens/models/qwen/qwen_model.h +++ b/SuperKittens/models/qwen/qwen_model.h @@ -395,7 +395,9 @@ inline void encode_gemm( enc->setBytes(&transA, 4, 9); enc->setBytes(&transB, 4, 10); enc->setBytes(&has_bias, 4, 11); enc->setBuffer(C, 0, 12); - enc->dispatchThreadgroups(MTL::Size((N + 63) / 64, (M + 63) / 64, 1), + // gemm_fp16 tiles BM=32 rows; a /64 row grid skips rows 32..63 of each + // 64-block at M>32 (same bug class fixed in deepseek_model.h). + enc->dispatchThreadgroups(MTL::Size((N + 63) / 64, (M + 31) / 32, 1), MTL::Size(64, 1, 1)); } From 1670c4d77aa50e52c30565c458d97a8618fcfb25 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 15:47:25 -0400 Subject: [PATCH 03/31] =?UTF-8?q?phi4:=20Phi-4-reasoning=20(14B)=20adapter?= =?UTF-8?q?=20over=20shared=20dense=20core=20=E2=80=94=20config-only=20+?= =?UTF-8?q?=20one-time=20phi3=20GGUF=20repack?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit phi3 arch fits the core exactly (head_dim 128, full rotary per phi3.rope.dimension_count=128, plain NeoX rope theta=5e5, GQA 40/10, untied Q6_K head, no qkv bias / qk-norm / sliding window). The only structural gap is artifact-side: phi3 GGUFs fuse attn_qkv ([Q;K;V]) and gate_up (as ffn_up, 2*n_ff rows); repack_phi3_gguf.py splits them offline by row-range byte copy (K-quant rows independent -> bit-exact, no requant), keeping the loader and kernels untouched. Gates on amelia (M4 16GB, CLT-only runtime compile, colima resident): load+config from GGUF header PASS; greedy coherence PASS (fluent reasoning, correct unit math, finite logits); Qwen3-1.7B-Q8 32-tok A/B vs pristine origin/main TOKEN-IDENTICAL (also regression-checks the bf05198 fp16-GEMM row-grid fix); decode median 10.49 tok/s pure / 9.79 generate (non-canonical host); batched N=4 lockstep smoke PASS. --- SuperKittens/inference/registry.py | 21 ++ .../models/load/tokenizer/chat_templates.py | 44 ++++ .../models/load/tokenizer/tokenizer.py | 4 + SuperKittens/models/phi4/__init__.py | 3 + SuperKittens/models/phi4/phi4.py | 53 +++++ SuperKittens/models/phi4/repack_phi3_gguf.py | 189 ++++++++++++++++++ temp/phi4_port/STATUS.md | 110 ++++++++++ 7 files changed, 424 insertions(+) create mode 100644 SuperKittens/models/phi4/__init__.py create mode 100644 SuperKittens/models/phi4/phi4.py create mode 100644 SuperKittens/models/phi4/repack_phi3_gguf.py create mode 100644 temp/phi4_port/STATUS.md diff --git a/SuperKittens/inference/registry.py b/SuperKittens/inference/registry.py index 0faa9c9..c588547 100644 --- a/SuperKittens/inference/registry.py +++ b/SuperKittens/inference/registry.py @@ -264,6 +264,27 @@ class ModelSpec: eps=1e-6, rope_freq_base=5_000_000.0, tie_word_embeddings=0, use_qk_norm=0), ), + # Phi-4-reasoning (14B): model_type "phi3" (Phi3ForCausalLM) — a dense decoder + # the shared core drives config-only: use_qk_norm=0, rope_interleaved=0 (phi3 + # GGUFs are NOT q/k-permuted → NeoX/type-2, the core default), plain RoPE + # theta=500000 (rope_scaling null, partial_rotary_factor 1.0 = full 128-dim), + # untied LM head, no QKV bias, no sliding window. vocab 100352 (tiktoken-style). + # ARTIFACT: phi3 GGUFs fuse attn_qkv + gate_up(ffn_up); gguf_name points at the + # one-time repack (models/phi4/repack_phi3_gguf.py — bit-exact row split). + # Q4_K_M (~8.4 GiB) fits a 16 GB mini with clamped cache_max. + "phi4-reasoning": ModelSpec( + family="phi4", + adapter="SuperKittens.models.phi4.phi4:Phi4", + hf_repo="microsoft/Phi-4-reasoning", + weight_dir="Phi-4-reasoning-GGUF", + gguf_name="microsoft_Phi-4-reasoning-Q4_K_M-sk.gguf", + default_quant="q4_k_m", + tokenizer_family="phi4", + dims=dict(n_layers=40, d_model=5120, n_heads=40, n_kv_heads=10, + head_dim=128, n_int=17920, vocab_size=100352, + eps=1e-5, rope_freq_base=500000.0, tie_word_embeddings=0, + use_qk_norm=0, rope_interleaved=0), + ), # Llama-3.2-1B-Instruct: same Llama arch as 3B but head_dim=64 (the 3B and # all other dense families are head_dim=128). Exercises the head_dim-templated # decode/causal attention (mha_*_64). Tied LM head, llama3 RoPE scaling. diff --git a/SuperKittens/models/load/tokenizer/chat_templates.py b/SuperKittens/models/load/tokenizer/chat_templates.py index c4ee089..7a7846a 100644 --- a/SuperKittens/models/load/tokenizer/chat_templates.py +++ b/SuperKittens/models/load/tokenizer/chat_templates.py @@ -99,6 +99,49 @@ def gemma4_template(messages: Sequence[dict], add_generation_prompt: bool = True return "".join(parts) +PHI4_REASONING_SYSTEM = ( + "You are Phi, a language model trained by Microsoft to help users. Your role " + "as an assistant involves thoroughly exploring questions through a systematic " + "thinking process before providing the final precise and accurate solutions. " + "This requires engaging in a comprehensive cycle of analysis, summarizing, " + "exploration, reassessment, reflection, backtracing, and iteration to develop " + "well-considered thinking process. Please structure your response into two " + "main sections: Thought and Solution using the specified format: " + "{Thought section} {Solution section}. In the Thought section, detail " + "your reasoning process in steps. Each step should include detailed " + "considerations such as analysing questions, summarizing relevant findings, " + "brainstorming new ideas, verifying the accuracy of the current steps, refining " + "any errors, and revisiting previous steps. In the Solution section, based on " + "various attempts, explorations, and reflections from the Thought section, " + "systematically present the final solution that you deem correct. The Solution " + "section should be logical, accurate, and concise and detail necessary steps " + "needed to reach the conclusion. Now, try to solve the following question " + "through the above guidelines:" +) + + +def phi4_template(messages: Sequence[dict], add_generation_prompt: bool = True) -> str: + """Phi-4(-reasoning) ChatML-with-<|im_sep|> format (no newlines between turns). + The official template hardcodes the reasoning system preamble; an explicit + leading system message in `messages` replaces it. No BOS is prepended. + """ + system = PHI4_REASONING_SYSTEM + body = [] + for m in messages: + role = m.get("role", "user").lower() + content = m.get("content", "") + if role == "system": + system = content + elif role == "user": + body.append(f"<|im_start|>user<|im_sep|>{content}<|im_end|>") + else: + body.append(f"<|im_start|>assistant<|im_sep|>{content}<|im_end|>") + parts = [f"<|im_start|>system<|im_sep|>{system}<|im_end|>"] + body + if add_generation_prompt: + parts.append("<|im_start|>assistant<|im_sep|>") + return "".join(parts) + + CHAT_TEMPLATES = { "deepseek": deepseek_template, "ds4": deepseek_template, @@ -111,4 +154,5 @@ def gemma4_template(messages: Sequence[dict], add_generation_prompt: bool = True "nemotron": llama_template, "mistral": mistral_template, "yi": qwen_template, + "phi4": phi4_template, } diff --git a/SuperKittens/models/load/tokenizer/tokenizer.py b/SuperKittens/models/load/tokenizer/tokenizer.py index beba8ab..e055d63 100644 --- a/SuperKittens/models/load/tokenizer/tokenizer.py +++ b/SuperKittens/models/load/tokenizer/tokenizer.py @@ -52,6 +52,10 @@ class Tokenizer: # bos is not prepended (chat(bos=False)). Distinct 64k vocab from Qwen. "yi": {"bos": ("<|startoftext|>",), "eos": ("<|im_end|>", "<|endoftext|>"), "pad": ("",)}, "mistral": {"bos": ("",), "eos": ("",), "pad": ("",)}, + # Phi-4(-reasoning): GPT-lineage tiktoken-style vocab. Chat turns end on + # <|im_end|> (100265); <|endoftext|> (100257) doubles as BOS and the + # base-completion stop. Pad is the repurposed <|dummy_85|> (100349). + "phi4": {"bos": ("<|endoftext|>",), "eos": ("<|im_end|>", "<|endoftext|>"), "pad": ("<|dummy_85|>",)}, # DeepSeek V2/V3 use full-width tokens for BOS/EOS in the trained vocab. "deepseek": {"bos": ("<｜begin▁of▁sentence｜>",), "eos": ("<｜end▁of▁sentence｜>",), diff --git a/SuperKittens/models/phi4/__init__.py b/SuperKittens/models/phi4/__init__.py new file mode 100644 index 0000000..d266b05 --- /dev/null +++ b/SuperKittens/models/phi4/__init__.py @@ -0,0 +1,3 @@ +from .phi4 import Phi4 + +__all__ = ["Phi4"] diff --git a/SuperKittens/models/phi4/phi4.py b/SuperKittens/models/phi4/phi4.py new file mode 100644 index 0000000..4cd4a55 --- /dev/null +++ b/SuperKittens/models/phi4/phi4.py @@ -0,0 +1,53 @@ +"""phi4.py — Phi-4-reasoning (14B) adapter over the dense core. + +Phi-4-reasoning is ``model_type=phi3`` (Phi3ForCausalLM) — a dense decoder the +shared :class:`DenseDecoder` core already drives. Config (verified against the +HF config.json AND the GGUF header): 40 layers, d_model 5120, 40 heads / 10 KV +heads, head_dim 128, n_int 17920, vocab 100352, rms_eps 1e-5. Configured: + + 1. ``use_qk_norm=0`` — Phi-3/4 has no per-head Q/K RMSNorm. + 2. ``rope_interleaved=0`` — phi3-arch GGUFs are NOT q/k-permuted (only the + llama arch gets the type-0 conversion permute), so the GGUF wants GGML + rope type 2 (NeoX/split-half), the core's default kernel. This matches HF + ``rotate_half`` on unpermuted weights — same convention as Qwen3. + 3. Plain RoPE, ``theta=500000``, ``rope_scaling=null``, + ``partial_rotary_factor=1.0`` — full 128-dim rotary, verified via the + GGUF's ``phi3.rope.dimension_count=128`` (a partial factor would need a + kernel change; the core has no support for it). The base + :meth:`DenseDecoder.bake_and_set_rope` is correct as-is. + 4. Untied LM head (``output.weight``) — ``tie_word_embeddings=0``. + 5. ``attention_bias=false`` — no QKV bias (no bias_add path). + 6. No BOS prepend: the tiktoken-style vocab is GPT-lineage; completion is not + BOS-sensitive and the chat template starts at ``<|im_start|>``. + +ARTIFACT NOTE: phi3-arch GGUFs ship FUSED projections — ``blk.N.attn_qkv`` +([Q;K;V] row-concat) and ``blk.N.ffn_up`` holding [gate;up] (2*n_ff rows). +The shared loader wants separate attn_q/k/v + ffn_gate/ffn_up tensors, so the +GGUF is repacked ONCE offline with ``repack_phi3_gguf.py`` (pure row-range byte +split — K-quant rows are independent, so the split is bit-exact, no requant). +The registry's ``gguf_name`` points at the repacked ``*-sk.gguf``. + +Usage: + import SuperKittens as sk + m = sk.load("phi4-reasoning") + print(m.chat("Why is the sky blue?")) +""" +from __future__ import annotations + +from SuperKittens.models.dense.dense_decoder import DenseDecoder + + +class Phi4(DenseDecoder): + """Phi-4-reasoning handle (own adapter; shared dense core).""" + + @classmethod + def from_spec(cls, spec, **overrides) -> "Phi4": + """Build from a registry ModelSpec. + + Uses the shared :meth:`DenseDecoder.from_spec` for config/GGUF/tokenizer, + forcing ``use_qk_norm=0`` (phi3 arch) and ``rope_interleaved=0`` (NeoX, + unpermuted GGUF). RoPE is plain (theta=5e5); the base bake is correct. + """ + overrides.setdefault("use_qk_norm", 0) + overrides.setdefault("rope_interleaved", 0) + return super().from_spec(spec, **overrides) diff --git a/SuperKittens/models/phi4/repack_phi3_gguf.py b/SuperKittens/models/phi4/repack_phi3_gguf.py new file mode 100644 index 0000000..bc05fec --- /dev/null +++ b/SuperKittens/models/phi4/repack_phi3_gguf.py @@ -0,0 +1,189 @@ +"""repack_phi3_gguf.py — split phi3-arch fused GGUF tensors for the SK loader. + +phi3-arch GGUFs (Phi-3/Phi-4 family) carry two fused per-layer projections: + + * ``blk.N.attn_qkv.weight`` — [Q;K;V] row-concatenated + (rows = n_head*head_dim + 2*n_kv_head*head_dim, ne0 = d_model) + * ``blk.N.ffn_up.weight`` — [gate;up] row-concatenated (rows = 2*n_ff) + +The shared dense loader (``sk_qwen_load_gguf``) expects the separate +qwen/llama-style tensors (attn_q/attn_k/attn_v, ffn_gate/ffn_up). GGUF tensor +data is row-major with each row quantized independently (K-quant superblocks +run along ne0), so splitting along the row dimension is a pure contiguous byte +copy — bit-exact, no dequant/requant. This tool rewrites the GGUF once, +offline; metadata KVs are copied verbatim (the SK loader takes dims from the +Python-side Config, not from GGUF KVs). + +Usage: + python3 repack_phi3_gguf.py in.gguf out.gguf + +Reads fused-split geometry (head counts, head_dim, n_ff) from the GGUF header. +Non-fused tensors are passed through untouched. +""" +from __future__ import annotations + +import struct +import sys + +ALIGN = 32 # GGUF default; general.alignment override is honored below. + +# GGML dtype code -> (block_size, bytes_per_block) for row-size math. +GGML_BLOCK = { + 0: (1, 4), # F32 + 1: (1, 2), # F16 + 2: (32, 18), # Q4_0 + 3: (32, 20), # Q4_1 + 6: (32, 22), # Q5_0 + 7: (32, 24), # Q5_1 + 8: (32, 34), # Q8_0 + 10: (256, 84), # Q2_K + 11: (256, 110), # Q3_K + 12: (256, 144), # Q4_K + 13: (256, 176), # Q5_K + 14: (256, 210), # Q6_K + 30: (1, 2), # BF16 +} + +_SCALAR_FMT = {0: "B", 1: "b", 2: "H", 3: "h", 4: "I", 5: "i", 6: "f", + 7: "?", 10: "Q", 11: "q", 12: "d"} + + +def _read_str(f): + (n,) = struct.unpack("" + else: + raise ValueError(f"bad kv type {vtype} for {key}") + + +def row_bytes(dtype: int, ne0: int) -> int: + bs, bb = GGML_BLOCK[dtype] + if ne0 % bs: + raise ValueError(f"ne0 {ne0} not divisible by block {bs} (dtype {dtype})") + return ne0 // bs * bb + + +def parse_header(path: str): + """Returns (kvs, tensors, kv_raw, data_start). tensors: [name, dims, dtype, off].""" + with open(path, "rb") as f: + magic = f.read(4) + if magic != b"GGUF": + raise ValueError(f"not a GGUF file: {magic!r}") + (version,) = struct.unpack(" None: + kvs, tensors, kv_raw, data_start, align, n_kv = parse_header(src) + + arch = kvs.get("general.architecture") + if arch != "phi3": + raise SystemExit(f"expected phi3 arch, got {arch!r}") + n_head = int(kvs["phi3.attention.head_count"]) + n_kv_head = int(kvs["phi3.attention.head_count_kv"]) + d_model = int(kvs["phi3.embedding_length"]) + n_ff = int(kvs["phi3.feed_forward_length"]) + head_dim = d_model // n_head + nq, nkv = n_head * head_dim, n_kv_head * head_dim + + out_tensors = [] # (name, dims, dtype, src_byte_off, nbytes) + for name, dims, dtype, off in tensors: + rb = row_bytes(dtype, dims[0]) + nrows = dims[1] if len(dims) > 1 else 1 + base = data_start + off + if name.endswith(".attn_qkv.weight"): + if nrows != nq + 2 * nkv: + raise SystemExit(f"{name}: rows {nrows} != q+2kv {nq + 2 * nkv}") + blk = name[: -len("attn_qkv.weight")] + out_tensors.append((blk + "attn_q.weight", [dims[0], nq], dtype, base, nq * rb)) + out_tensors.append((blk + "attn_k.weight", [dims[0], nkv], dtype, base + nq * rb, nkv * rb)) + out_tensors.append((blk + "attn_v.weight", [dims[0], nkv], dtype, base + (nq + nkv) * rb, nkv * rb)) + elif name.endswith(".ffn_up.weight") and nrows == 2 * n_ff: + blk = name[: -len("ffn_up.weight")] + out_tensors.append((blk + "ffn_gate.weight", [dims[0], n_ff], dtype, base, n_ff * rb)) + out_tensors.append((blk + "ffn_up.weight", [dims[0], n_ff], dtype, base + n_ff * rb, n_ff * rb)) + else: + out_tensors.append((name, dims, dtype, base, nrows * rb)) + + # New tensor-info block with recomputed (aligned) data offsets. + infos = bytearray() + data_off = 0 + placed = [] # (src_off, nbytes, dst_rel_off) + for name, dims, dtype, src_off, nbytes in out_tensors: + data_off = (data_off + align - 1) // align * align + nb = name.encode("utf-8") + infos += struct.pack(" {len(out_tensors)} tensors; " + f"geometry: heads={n_head} kv={n_kv_head} head_dim={head_dim} n_ff={n_ff}") + + +if __name__ == "__main__": + if len(sys.argv) != 3: + raise SystemExit(__doc__) + repack(sys.argv[1], sys.argv[2]) diff --git a/temp/phi4_port/STATUS.md b/temp/phi4_port/STATUS.md new file mode 100644 index 0000000..4ab55ba --- /dev/null +++ b/temp/phi4_port/STATUS.md @@ -0,0 +1,110 @@ +# Phi-4-reasoning (14B) port — STATUS + +**Verdict: PORTED-COHERENT** (config-only adapter over the shared dense core + +a one-time offline GGUF repack; ZERO kernel / launcher / loader edits). + +Branch: `dev-sk-phi4` (based on local `main` @ bf05198, which includes the +fp16-GEMM row-grid fix — gate 3 below doubles as its regression check). +Bench host: amelia (M4 mini 16 GB, CLT-only, colima resident — NOT the +canonical lexie protocol). + +## Config (verified from GGUF header + HF config.json, not memory) + +| field | value | source | +|---|---|---| +| arch | `phi3` (Phi3ForCausalLM) | `general.architecture` | +| n_layers | 40 | `phi3.block_count` | +| d_model | 5120 | `phi3.embedding_length` | +| n_heads / n_kv_heads | 40 / 10 (GQA 4:1) | `phi3.attention.head_count{,_kv}` | +| head_dim | 128 (= core's only supported dim) | 5120/40; `phi3.rope.dimension_count` | +| partial rotary | **NONE** — `rope.dimension_count=128` = full | GGUF header | +| rope | plain NeoX (type-2), theta=500000, `rope_scaling=null` | GGUF + config.json | +| rope_interleaved | 0 (phi3 GGUFs are NOT q/k-permuted; only llama-arch gets the type-0 permute) | llama.cpp convert behavior | +| n_int (FFN) | 17920 | `phi3.feed_forward_length` | +| vocab | 100352 (tiktoken-style GPT-lineage) | token_embd dims | +| eps | 1e-5 | `phi3.attention.layer_norm_rms_epsilon` | +| qkv bias / qk-norm / sliding window | none / none / 0 | config.json + GGUF | +| LM head | UNTIED `output.weight`, Q6_K → native q6k_matvec | GGUF | +| BOS / EOS / pad | 100257 `<|endoftext|>` / 100265 `<|im_end|>` (+100257 stop) / 100349 `<|dummy_85|>` | tokenizer_config.json | +| chat format | `<|im_start|>role<|im_sep|>…<|im_end|>` + hardcoded reasoning system preamble | tokenizer_config chat_template | +| BOS-prepend | not needed (GPT-lineage; template starts at `<|im_start|>`) | — | + +Quant mix (bartowski Q4_K_M, imatrix): attn_qkv Q5_K (uniform), attn_output +Q4_K, ffn gate/up Q4_K, ffn_down Q4_K/Q6_K per-layer mix, embed Q4_K +(host-dequant fp16), head Q6_K. All dtypes already supported +(Q5_K is an in-tree fit-enabler). + +## The one structural wrinkle: fused GGUF tensors → offline repack + +phi3-arch GGUFs fuse `blk.N.attn_qkv.weight` ([Q;K;V] row-concat) and +`blk.N.ffn_up.weight` ([gate;up], 2*n_ff rows). The shared loader +(`sk_qwen_load_gguf`) wants separate `attn_q/k/v` + `ffn_gate/ffn_up`. +Rather than teach the core a fused path (out of scope for a breadth port), +`SuperKittens/models/phi4/repack_phi3_gguf.py` rewrites the GGUF once: +pure row-range byte split (K-quant rows are independent → bit-exact, no +dequant/requant; HF order is q,k,v and gate,up). 243 → 363 tensors, ++7 KB file size. Validated: md5 byte-identity of every split range for +layers 0/20/39 + embed/head/down spot checks, and `gguf` pkg re-parse +(363 tensors, all expected names present) — PASSED. + +Artifact on amelia: `~/phi4-gguf/microsoft_Phi-4-reasoning-Q4_K_M-sk.gguf` +(9.05 GB). The fused original was deleted for disk headroom; regenerate via +`curl -L` of `bartowski/microsoft_Phi-4-reasoning-GGUF` Q4_K_M + one repack run +(~3 min total). Sidecar files in the same dir: config.json, tokenizer.json, +tokenizer_config.json (from `microsoft/Phi-4-reasoning`). + +## What was added (all Python, zero core edits) + +- `SuperKittens/models/phi4/phi4.py` — `Phi4(DenseDecoder)`, the Mistral/Yi + pattern: `use_qk_norm=0`, `rope_interleaved=0`; base RoPE bake correct as-is. +- `SuperKittens/models/phi4/repack_phi3_gguf.py` — one-time GGUF repack tool. +- `inference/registry.py` — `"phi4-reasoning"` ModelSpec row. +- `models/load/tokenizer/tokenizer.py` — `"phi4"` family specials row. +- `models/load/tokenizer/chat_templates.py` — `phi4_template` + (`<|im_sep|>` ChatML variant + official reasoning system preamble). + +## Gates + +1. **Load + config print — PASS.** Loads the repacked GGUF in 13.4s; printed + config matches the table above; tokenizer resolves + bos=100257 / eos=100265 / eos_ids={100257,100265} / pad=100349. +2. **Coherence (greedy) — PASS.** Finite logits on all runs. + - raw "Generate a poem about pizza dough:" (64 tok): fluent constraint-style + continuation (instruct model completing without template — expected style). + - chat-templated poem (96 tok): fluent `` planning scaffold. + - chat-templated reasoning ("train 60 miles in 40 min, how far in 2h?", + 96 tok): correct reasoning — converts units, derives 1.5 mi/min ✓. +3. **No-regression A/B (Qwen3-1.7B-Q8, 32-tok greedy, same prompt) — PASS.** + dev-sk-phi4 build (bf05198+adapter) vs pristine origin/main (3491001) + build: token-identical (see `lab_qwen_{A,B}.log`). Also covers the + row-grid fix regression check. +4. **Decode tok/s — RECORDED** (2 warmup + 5 reps, median, 0.3 s gaps, + caffeinate -is; amelia-with-colima, NOT the canonical lexie protocol): + - pure decode (per-token forward after prefill): **10.49 tok/s** + (reps 10.46-10.51 — thermally tight) + - generate-loop (bench.py convention, ~8-tok prompt + 64 new): 9.79 tok/s + - first cold run pages in the 9 GB mmap: 4.58 tok/s (warmup only) + - context: qwen3-14B Q4_K_M is 11.28 tok/s on lexie; phi4 has a slightly + larger FFN (17920 vs 17408) and this host carries colima. +5. **Batched N=4 lockstep smoke — PASS** (bonus; `batch=4, cache_max=256`, + `prefill_batched` + 8 `forward_batched` steps, distinct prompts): all ids + in vocab range, no NaN, all four lanes coherent and prompt-relevant + ("…100°C at sea level", "…concise way to create lists", …). Minor word + echo at lane starts is the lockstep equal-length prompt trim, not the model. + +## Memory notes (16 GB + colima) + +- Pinned `cache_max=512, seq_max=512` (single-stream gates); ONE handle per + process, no sequential reloads. +- Swap during the 14B runs: 1.18 G → 2.18 G used (total grew to 3 G); below the + 3.5 G back-off line, stable across bench reps and the batched run. + +## How to run (amelia) + +```sh +cd ~/sk-phi4-port && source skenv.sh # SK_DYLIB + SK_METAL_SRC_FALLBACK + PYTHONPATH +python3 lab_phi4_gates12.py # load + coherence +python3 lab_phi4_bench.py # decode tok/s +# A/B: SK_ROOT=$HOME/sk-phi4-port/base_main SK_DYLIB=$PWD/build/libsk_base.dylib \ +# source skenv.sh && python3 lab_qwen_identity.py +``` From c75ea97f5793ff529559683bce999a76dbebe694 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 15:57:01 -0400 Subject: [PATCH 04/31] =?UTF-8?q?gemma4:=20fix=20fp16-GEMM=20row=20grid=20?= =?UTF-8?q?at=204=20sites=20(qkv=5Fpacked/o=5Fproj=20fallbacks=20+=202=20P?= =?UTF-8?q?LE-inject=20GEMMs)=20=E2=80=94=20gemm=5Ffp16=20tiles=20BM=3D32,?= =?UTF-8?q?=20/64=20grid=20skipped=20rows=20at=20M>32=20(audit=20per=20dee?= =?UTF-8?q?pseek=20bug=20class)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SuperKittens/models/gemma/gemma4/gemma4_model.h | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/SuperKittens/models/gemma/gemma4/gemma4_model.h b/SuperKittens/models/gemma/gemma4/gemma4_model.h index 8fbc989..ab8c295 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4_model.h +++ b/SuperKittens/models/gemma/gemma4/gemma4_model.h @@ -429,7 +429,7 @@ inline void dispatch_layer( enc->setBytes(&transB, 4, 10); enc->setBytes(&has_bias, 4, 11); enc->setBuffer(B.qkv_packed, 0, 12); - enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 63) / 64, 1), + enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 31) / 32, 1), MTL::Size(64, 1, 1)); } } @@ -765,7 +765,7 @@ inline void dispatch_layer( enc->setBytes(&transB, 4, 10); enc->setBytes(&has_bias, 4, 11); enc->setBuffer(B.o_proj, 0, 12); - enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 63) / 64, 1), + enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 31) / 32, 1), MTL::Size(64, 1, 1)); } } @@ -1036,7 +1036,7 @@ inline void dispatch_ple_inject( enc->setBytes(&transB, 4, 10); enc->setBytes(&has_bias, 4, 11); enc->setBuffer(B.ple_gate_out, 0, 12); - enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 63) / 64, 1), + enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 31) / 32, 1), MTL::Size(64, 1, 1)); enc->endEncoding(); } @@ -1084,7 +1084,7 @@ inline void dispatch_ple_inject( enc->setBytes(&transB, 4, 10); enc->setBytes(&has_bias, 4, 11); enc->setBuffer(B.ple_proj_back, 0, 12); - enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 63) / 64, 1), + enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M + 31) / 32, 1), MTL::Size(64, 1, 1)); enc->endEncoding(); } @@ -1454,7 +1454,7 @@ inline void dispatch_model( enc2->setBytes(&transB, 4, 10); enc2->setBytes(&has_bias,4, 11); enc2->setBuffer(B.ple_ctx_proj, 0, 12); - enc2->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M_v + 63) / 64, 1), + enc2->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M_v + 31) / 32, 1), MTL::Size(64, 1, 1)); enc2->endEncoding(); } @@ -1942,7 +1942,7 @@ inline void dispatch_model( enc->setBytes(&transB, 4, 10); enc->setBytes(&has_bias, 4, 11); enc->setBuffer(B.logits, 0, 12); - enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M_v + 63) / 64, 1), + enc->dispatchThreadgroups(MTL::Size((N_v + 63) / 64, (M_v + 31) / 32, 1), MTL::Size(64, 1, 1)); enc->endEncoding(); } From 52b332c59d0de3f3c9a6e67c84da9a4b165610b6 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 16:04:21 -0400 Subject: [PATCH 05/31] =?UTF-8?q?gemma4=20batched=20prefill:=20VERIFIED=20?= =?UTF-8?q?WIN=20=E2=80=94=20lane=20gates=20closed,=20divergence=20root-ca?= =?UTF-8?q?used=20as=20inherited=20seq>1-vs-seq=3D1=20attention=20numerics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit STATUS + gate drivers + artifacts. Headline: N=8 T=128 serving TTFT 5.11x (single chunk) / 5.54x (chunk=64) on derek M4 16GB, decode aggregate 1.0000, old paths byte-identical (final build, batch=1 + batch=8). Lane correctness: chunk-split==single-chunk 8/8; identical-prompt and bitwise lane-permutation invariants hold on random/raw-fluent/chat prompt sets; first token 8/8 vs BOTH references (token-by-token lockstep AND single-stream forward(T)) in all regimes; 32-tok continuations 7/8 vs the single-stream seq>1 reference across different decode engines. The non-8/8 continuation-vs-token-by-token gate (6/8 chat, 3/8 raw, 2/8 random) is pre-existing: zero-new-code base-dylib demo shows the EXISTING single-stream prefill diverges from token-by-token at the same rates on the same prompts (6/8 chat sharing prompt 1; 4/8 raw, one at index 0) — the q_seq==1 split-Bc decode fast path vs q_seq>1 sequential online-softmax orderings differ in bf16 rounding. Earlier 'degenerate baseline' mystery was the harness tiling prompts x10. Row-grid audit table in STATUS: 6 BUG sites (committed in c75ea97), all dormant for 12B-unified (no PLE, Q4K body, Q8 embed head); latent repairs for PLE E-variants + bf16 fallbacks. --- temp/gemma4_bprefill/STATUS.md | 266 +++++++ .../gemma4_bprefill/artifacts/gate_chat2.json | 603 ++++++++++++++++ .../artifacts/gate_fluent_r4.json | 35 + .../artifacts/gates_12b_fixed.json | 675 ++++++++++++++++++ .../artifacts/seq_paths_chat.json | 586 +++++++++++++++ .../artifacts/seq_paths_fluent.json | 586 +++++++++++++++ temp/gemma4_bprefill/artifacts/ss_new_r4.json | 73 ++ temp/gemma4_bprefill/gate_chat.py | 107 +++ temp/gemma4_bprefill/gate_fluent.py | 106 +++ temp/gemma4_bprefill/gate_real_prompts.py | 99 +++ temp/gemma4_bprefill/gates_ttft.py | 157 ++++ temp/gemma4_bprefill/lockstep_identity.py | 44 ++ temp/gemma4_bprefill/prompts_fluent.py | 10 + temp/gemma4_bprefill/seq_paths_base.py | 58 ++ temp/gemma4_bprefill/seq_paths_chat.py | 63 ++ temp/gemma4_bprefill/seq_paths_fluent.py | 52 ++ temp/gemma4_bprefill/single_stream.py | 47 ++ 17 files changed, 3567 insertions(+) create mode 100644 temp/gemma4_bprefill/STATUS.md create mode 100644 temp/gemma4_bprefill/artifacts/gate_chat2.json create mode 100644 temp/gemma4_bprefill/artifacts/gate_fluent_r4.json create mode 100644 temp/gemma4_bprefill/artifacts/gates_12b_fixed.json create mode 100644 temp/gemma4_bprefill/artifacts/seq_paths_chat.json create mode 100644 temp/gemma4_bprefill/artifacts/seq_paths_fluent.json create mode 100644 temp/gemma4_bprefill/artifacts/ss_new_r4.json create mode 100644 temp/gemma4_bprefill/gate_chat.py create mode 100644 temp/gemma4_bprefill/gate_fluent.py create mode 100644 temp/gemma4_bprefill/gate_real_prompts.py create mode 100644 temp/gemma4_bprefill/gates_ttft.py create mode 100644 temp/gemma4_bprefill/lockstep_identity.py create mode 100644 temp/gemma4_bprefill/prompts_fluent.py create mode 100644 temp/gemma4_bprefill/seq_paths_base.py create mode 100644 temp/gemma4_bprefill/seq_paths_chat.py create mode 100644 temp/gemma4_bprefill/seq_paths_fluent.py create mode 100644 temp/gemma4_bprefill/single_stream.py diff --git a/temp/gemma4_bprefill/STATUS.md b/temp/gemma4_bprefill/STATUS.md new file mode 100644 index 0000000..17830c6 --- /dev/null +++ b/temp/gemma4_bprefill/STATUS.md @@ -0,0 +1,266 @@ +# Batched (N-lane) seq>1 chunked prefill — gemma4_unified (12B) + +Branch: `dev-sk-gemma-bprefill` (off local main @71df9ea, which carries the +qwen batched prefill). Bench host: derek (M4 base, 16 GB, CLT-only; dylib +clang++-built, kernels runtime-compiled via SK_METAL_SRC_FALLBACK; colima +default+k3s VMs resident throughout). Model: gemma-4-12B-it Q4_K_M GGUF, +SK_GEMMA4_BODY_Q4K=1 SK_GEMMA4_EMBED_Q8=1 (the PR #78 fit-16GB knobs). + +## Problem +`sk_gemma4_forward_batched` rejects seq!=1 (-6): serving callers ingest +prompts token-by-token at seq==1 (weight-amortized across lanes but paying T +full weight-read passes for a T-token prompt). The qwen fix (PR'd earlier) +does one M=batch*chunk GEMM pass per chunk instead. + +## Audit findings (what gemma4 needed vs qwen) +The qwen breakage was lane-0-only seq<->head transposes. gemma4 has NO such +transposes, and the batch-aware kernels added for lockstep decode (667c2eb) +turn out to be seq-generic already: + +- `gemma4_qkv_norm_batched` maps flat row t -> lane b=t/seq, pos s=t%seq and + writes (B,H,seq,D) — any seq. +- `rope_qk_bf16_batched` / `gemma4_rope_qk_partial_batched` rotate row r at + write_pos + (r%seq) over the (B,H,seq,D) layout — any seq. +- `kv_cache_write_bf16` reads (B,H_kv,seq,D), writes per-lane ring slices + (B,H_kv,csize,D) at (pos+t)%csize — any seq, any batch. +- `gemma4_attn_local_d256`/`gemma4_attn_global_d512` index Q at + (batch*nheads+head)*seq*D, per-lane KV slice at + (batch*n_kv_heads+kv_head)*csize*D, grid z=batch; the seq>1 "original path" + handles causal masking per query row. O is written (B,seq,H*D) = flat + request-major rows for the o_proj GEMM. +- Projections/norms/MLP/layer_scalar operate on flat T=batch*seq rows + (M-agnostic `gemma4_gemm_mma_*` body GEMM at M>1). + +What was actually missing (all additive, default-off): +1. A tail that doesn't project all batch*seq rows: the existing T>1 head + either GEMMs all T rows (bf16 embed) or loops the quant matvec per row — + vocab(262144) x T dead work — and `decode_all_rows` argmax requires seq==1. +2. A chunk-loop ABI carrying positions/KV across chunks. + +### SWA audit (the flagged silent-corruption spot) +Per-lane KV slices + ring addressing are batch-correct at seq>1 (above). The +real SWA finding is an EXACTNESS BOUND inherited from the single-stream seq>1 +path, not a batch bug: + +- Local layers use a ring buffer of size `window` (csize==window), so the + in-kernel window mask (`lower = upper>window ? upper-window : 0`) is dead + code there — kv_len <= window always; SWA is enforced by ring EVICTION. +- When a whole chunk is written before attention reads (the dispatch order), + the chunk's tail overwrites the oldest in-window keys of its own interior + rows once current_pos+seq > window: query row r of the chunk loses + (seq-1-r) of its oldest window keys. The mask accounts for exactly the + retained keys (no garbage reads) — outputs are a slightly-shorter-window + approximation, diverging from the token-by-token reference. +- Bound: chunked seq>1 prefill is token-EXACT iff current_pos+seq <= window + (then nothing is evicted: kv_len = total tokens, lower=0, full causal + visibility). Token-by-token (seq=1) is exact at ANY length — each step's + single query row is the chunk's last row. +- This is a pre-existing property of `sk_gemma4_forward` at seq>1 (same + write-then-attend order, same ring); batching does not widen it. Documented + on the new ABI rather than "fixed": an exact >window chunked prefill needs + attend-before-overwrite or scratch K/V — out of scope here. +- Global layers (every 6th, csize=cache_max) are exact for all T <= cache_max + (enforced by the -4 bounds check). + +## Row-grid audit (deepseek 31ff3c4 bug class) — 6 real sites, ALL DORMANT for 12B +`gemma4_model.h` had 8 dispatches dividing the row grid by 64. Audited each +against the kernel it actually binds (a /64 grid against a BM=32 kernel +silently skips rows 32-63 of every 64-block at M>32; M<=32 unaffected): + +| site (pre-fix line) | binds | kernel BM | grid.y was | verdict | +|---|---|---|---|---| +| L432 QKV proj fallback | `gemm_bf16` (kernels/gemm/bf16/gemm.metal) | 32 | (M+63)/64 | **BUG** — only when `!use_q8` (bf16-weight fallback); not exercised with BODY_Q4K | +| L768 o_proj fallback | `gemm_bf16` | 32 | (M+63)/64 | **BUG** — same fallback-only | +| L908 fused MLP | `gated_mlp_bf16` | 64 | (M+63)/64 | correct | +| L1039 PLE gate | `gemm_bf16` | 32 | (M+63)/64 | **BUG** at T>32 — but PLE only | +| L1087 PLE proj | `gemma4_gemm_bf16_fp32_out` (ple_inject.metal) | 32 | (M+63)/64 | **BUG** at T>32 — PLE only | +| L1457 PLE ctx proj | `gemm_bf16` | 32 | (M_v+63)/64 | **BUG** at T>32 — PLE only | +| L1839 batched lane head | `gemm_bf16` | 32 | 1 (M=1) | correct | +| L1945 LM head bf16 T>1 | `gemm_bf16` | 32 | (M_v+63)/64 | **BUG** at T>32 — only when quant head absent AND w_embed present | +| enc_body (all body projs) | `gemma4_gemm_mma_*_t64n` / `_t32` | 64 / 32 | (M+BM-1)/BM, BM matched to variant | correct | + +All six buggy sites fixed to `(M + 31) / 32` on this branch. **Decisive for +this model: NONE of them fire for gemma4-12b-unified** — `gemma4_unified.py` +sets `has_ple=False` (kills L1039/L1087/L1457), BODY_Q4K=1 routes projections +through the correctly-dispatched mma kernels (kills L432/L768), and +EMBED_Q8=1 frees w_embed so the T>1 head loops the quant matvec per row +(kills L1945). Verified empirically: the fixed build reproduces the pre-fix +gate-1 divergence pattern bit-for-bit (same 6/8 lanes, same indices) — the +12B lane divergence is NOT the row-grid bug. The fixes are kept as a latent +correctness repair for the PLE-bearing E-variants (E2B/E4B run T>32 prefill +through L1039/L1087/L1457 and silently lose PLE injection on rows 32+ of +every 64-row block) and for the bf16-weight fallbacks; M<=32 and decode +dispatch grids are unchanged by construction ((M+63)/64 == (M+31)/32 == 1). + +## Design (mirrors the qwen shape; existing paths byte-identical) +1. `ModelParams.batched_prefill` (default 0): 1 = interior chunk — skip final + norm + head + descale + softcap + argmax entirely; 2 = final chunk — + project ONLY each lane's last prompt row (b*seq+seq-1) -> logits row b + (per-lane matvec: Q4_K head via `q4k_matvec_bf16`, else Q8_0, else bf16 + GEMM M=1), then `gemma4_logit_descale` (1/sqrt(d_model)) + + `gemma4_logit_softcap` (cap=30) over the batch*vocab lane rows, then + per-lane argmax -> output_id[b]. Gemma's softcapped-logits semantics are + preserved bit-for-bit with the decode tail (same kernels, same shapes). +2. New ABI `sk_gemma4_prefill_batched(h, ids[batch*seq] request-major, seq, + chunk_size, out_next[batch])`: chunk loop (<= seq_max), positions derived + from current_pos (gemma4 has no rope_pos buffer), KV carried across chunks. +3. Python `Gemma4.prefill_batched(ids, chunk_size=0)` — symbol-gated. + +Bans respected: zero changes to kernels/gemm/gemm_mma.metal (the gemma body +GEMM is models/gemma/gemma4/gemma4_gemm_mma.metal, also untouched); decode and +M=1 paths byte-identical (batched_prefill defaults 0 everywhere). + +## Gates +All runs: derek, one handle per process, detached + polled; batch=8, +seq_max=128, cache_max=512, window OVERRIDDEN 1024->512 for memory headroom +(KV ring halves; with every total ≤ 512 tokens nothing is ever evicted, so +compute/tokens are identical to window=1024 — by the ring analysis above). + +1. **Baseline coherence — PASS.** Base dylib, batch=1, greedy 48 tokens on the + chat-templated "Generate a poem about pizza dough": coherent verse + ("A mound of flour, like fallen snow, / In a wooden bowl where the shadows + flow..."). +2. **Lane isolation — PASS with a documented numerical caveat (see below).** + N=8 distinct prompts, T=128, chunk=64: + - first token after batched prefill == token-by-token baseline: 8/8; + - chunk=64 == single-chunk next tokens: 8/8; + - identical prompts in all lanes -> identical outputs: yes; + - bitwise lane-permutation test (same code path both runs): permuting the + prompts permutes next tokens AND 8-token continuations exactly — no + cross-lane leakage, no lane-slot dependence (PASS on random-token, + raw-fluent, and chat-templated prompt sets); + - 32-token continuation vs the tbt baseline depends on argmax sharpness + (all on the post-row-grid-fix build): + * random-token prompts: first tokens 8/8, continuations 2/8 + (gates_12b_fixed.json, window=1024) — A-side streams are degenerate + newline/comma junk, knife-edge ties everywhere; + * RAW fluent passages (no chat template, T=120): first tokens 8/8, + continuations 3/8 (gate_fluent_r4.log) — the it-model greedy-parrots + untemplated text into repetition loops ('. Today laser altimetry. + Today.\n\n111...'), again knife-edge; + * CHAT-TEMPLATED fluent passages (sharp argmax, the it-model's regime; + full gemma chat turn "Continue this passage: ..." spliced to exactly + T=128 with the model-turn cue intact, gate_chat.py): first tokens + 8/8; A-side continuations genuinely fluent prose on all 8 lanes; + 32-tok continuations 6/8 — lanes 1 and 5 flip mid-continuation + (both sides fluent paraphrases of each other). Deterministic: a + second run (gate_chat2.json) reproduces every token. + - vs the SINGLE-STREAM seq>1 prefill reference (mission 3b; same q_seq>1 + attention path as the chunked prefill; seq_paths_chat.json c1 vs + gate_chat2.json): per-lane FIRST TOKEN 8/8 exact; 32-tok continuations + 7/8 exact EVEN THOUGH the continuation engines differ (batched M=8 + lockstep decode vs batch=1 M=1 decode) — the single flip is lane 1, the + same knife-edge prompt where the two EXISTING engines flip against each + other (demo below). The batched chunked prefill tracks its own + reference class (seq>1 prefill) more tightly than either tracks + token-by-token. + Root cause of the residual flips is PRE-EXISTING numerics, not the new + code: the A side prefills through the attention q_seq==1 decode fast + path (split-Bc + cross-simd merge), the B side through the q_seq>1 + original path — same online softmax, different reduction order, low-bit + bf16 noise (see the base-dylib demo below). The earlier "real-text" run + (gates_real.log, both sides DEGENERATE on diverging lanes) is explained: + gate_real_prompts.py tiled each sentence x10 to reach T=128, and the + 12B-it model legitimately degenerates on 10x-repeated text — harness + prompt construction, not a model or kernel bug. +3. **Old paths byte-identical — PASS.** Same host/env, base (branch-base + 71df9ea files) vs patched dylib: batch=1 chat coherence ids + greedy + forward+16-step decode ids token-identical; batch=8 lockstep tbt T=32 + + 16-step continuation token-identical. Held on ALL THREE builds: pre-fix + (r2 ls_base.json == ls_new.json), post-row-grid-fix (r3 + ls_base_r3.json == ls_new_r3.json), and the FINAL r4 build (ss_new_r4.json + coherence_ids + decode_ids == r2 ss_base.json) — expected by construction: + every dispatch the 12B config reaches has M<=32 at decode, where + (M+63)/64 == (M+31)/32 == 1. +4. **TTFT >= 8% — PASS** (80.4-81.9% improvement, table below). +5. **Decode aggregate after prefill — PASS:** 10.918 tok/s (after A) vs + 10.918 tok/s (after B), ratio 1.0000 (within +-2%) — post-fix build, + window=1024 (gates_12b_fixed.json). + +### Base-dylib inherited-divergence demo (zero new code) +seq_paths_fluent.py / seq_paths_chat.py, batch=1, BASE dylib (71df9ea files): +for the same prompts, path1 = one forward(T) (the EXISTING validated +single-stream seq>1 prefill, q_seq>1 attention) vs path2 = T x forward(seq=1) +(decode fast path), 32 greedy continuations each. Results: RAW fluent +prompts (T=120): 4/8 match — 4/8 diverge, one (prompt 5) at continuation +index 0, i.e. the two EXISTING paths disagree on the very first generated +token. Chat-templated (T=128): 6/8 match — prompts {1,7} diverge (idx 14 / +deeper). Compare the batched gate: 3/8 (raw) and 6/8 (chat, lanes {1,5}, +sharing prompt 1) — same regime-dependent rates, overlapping prompts. The +continuation-vs-tbt flip is a property of gemma4's seq>1-vs-seq=1 attention +numerics that predates and is untouched by the batched prefill — an 8/8 +32-tok-continuation-vs-tbt gate is unattainable for ANY seq>1 prefill on this +model, including the production single-stream one. + +SWA honesty (gate-2 clause): window=512 (overridden; production 1024) and +T=128 < window, so the SWA ring-wrap regime is UNEXERCISED by these gates — +deliberately: per the audit above, seq>1 chunked prefill beyond the window is +not token-exact BY CONSTRUCTION (pre-existing ring-eviction bound, documented +on the ABI), so an exercised-wrap equality gate would be vacuous. Callers must +keep context+prompt <= window for exact prefill (12B production window 1024). + +## A/B: serving TTFT, N=8 lanes (median of 7, 2 warmups, 0.3s gaps, same process/handle) +A = token-by-token lockstep prefill (`forward_batched` seq=1 x T); B = +`sk_gemma4_prefill_batched`. TTFT = wall from reset to all-lanes-first-token. +gemma-4-12B-it Q4_K_M, derek M4 base 16 GB, colima VMs resident. + +Final (post-row-grid-fix) build, window=1024 (production), gates_12b_fixed.json: + +| T | chunk | A (tbt) | B (batched) | speedup | improvement | +|---|---|---|---|---|---| +| 128 | 128 (1 chunk) | 93360.3 ms | 18275.0 ms | 5.11x | 80.4% | +| 128 | 64 | 93360.3 ms | 16856.9 ms | 5.54x | 81.9% | + +(The pre-fix r2 run at window=512 measured 5.13x/5.56x — the fix is +perf-neutral for 12B since none of the fixed sites fire in this config.) +T=256 not run: logits scratch is T_max*vocab(262144) bf16 — batch=8 +seq_max=256 would add ~1.07 GB on a box already near its ceiling with colima +VMs resident (qwen-8B precedent: headline scoped to T=128). + +chunk=64 measured ~7.6% faster than one 128-chunk (interior chunks attend to +shorter KV prefixes; the body GEMM is already weight-amortized at M=512). + +## Verdict +**WIN — promote.** Perf: serving TTFT 5.11x (single chunk) / 5.54x (chunk=64) +at N=8 T=128, decode aggregate ratio 1.0000, old paths byte-identical on the +final build. Correctness: every check that isolates the NEW code is exact — +chunk-split == single-chunk 8/8; identical prompts -> identical lanes; +lane-permutation bitwise-exact (3 prompt regimes); deterministic across +re-runs; first token 8/8 vs BOTH references (token-by-token AND single-stream +forward(T)) in every regime; 32-tok continuations 7/8 vs the single-stream +seq>1 reference even across different decode engines. + +The one gate that is NOT 8/8 — 32-tok continuations vs the token-by-token +reference (6/8 chat, 3/8 raw-fluent, 2/8 random) — is shown by the zero-new- +code base-dylib demo to be a pre-existing property of gemma4's seq>1-vs-seq=1 +attention numerics: the EXISTING production single-stream prefill diverges +from token-by-token at the same rates on the same prompts (6/8 chat sharing +prompt 1; 4/8 raw, one at continuation index 0). No seq>1 prefill on this +model can pass that gate as stated; the batched prefill is exactly as faithful +as the validated single-stream prefill it generalizes. No evidence of any +lane defect remains; the row-grid audit's six /64-vs-BM=32 fixes are kept as +latent repairs for the PLE E-variants and bf16-weight fallbacks (none fire in +the 12B config). + +NOT root-caused to a fixable site (and out of scope): which of the q_seq==1 +split-Bc merge vs q_seq>1 sequential softmax orderings is "righter" — they are +both valid online-softmax reductions; making them bit-identical would mean +rewriting the decode fast path as a degenerate seq>1 tile pass and paying its +latency. + +## Files +- `SuperKittens/models/gemma/gemma4/gemma4_model.h` — `ModelParams.batched_prefill`, + interior-chunk early-out, final-chunk per-lane head/descale/softcap/argmax tail. +- `SuperKittens/models/gemma/gemma4/launcher.c++` — `sk_gemma4_prefill_batched` + chunk loop. +- `SuperKittens/models/gemma/gemma4/launcher.h` — ABI decl. +- `SuperKittens/models/gemma/gemma4/gemma4.py` — `prefill_batched()` wrapper. +- `temp/gemma4_bprefill/{gates_ttft.py,single_stream.py,lockstep_identity.py}` — + gate/bench drivers; `derek/{build_dylib.sh,skenv.sh}` — CLT-only build + + runtime-compile env; `base_files/` — branch-base copies for the A/B dylib. +- `temp/gemma4_bprefill/{gate_chat.py,seq_paths_fluent.py,seq_paths_chat.py}` — + chat-templated lane gate (exact-T template splice) + zero-new-code + inherited-divergence demos; `artifacts/` — result JSONs (gate_chat2, + gates_12b_fixed, seq_paths_fluent/chat, ss identity). +- derek lab dirs: ~/sk-gemma-bprefill-r2 (pre-fix gates), -r3 (post-fix + gates/TTFT), -r4 (final build, fluent/chat gates, demos, ss identity). diff --git a/temp/gemma4_bprefill/artifacts/gate_chat2.json b/temp/gemma4_bprefill/artifacts/gate_chat2.json new file mode 100644 index 0000000..20ecf71 --- /dev/null +++ b/temp/gemma4_bprefill/artifacts/gate_chat2.json @@ -0,0 +1,603 @@ +{ + "first_match": [ + true, + true, + true, + true, + true, + true, + true, + true + ], + "cont_match": [ + true, + false, + true, + true, + true, + false, + true, + true + ], + "same_prompt_ok": true, + "perm_next_ok": true, + "perm_cont_ok": true, + "A_next": [ + 100, + 100, + 100, + 100, + 100, + 100, + 100, + 100 + ], + "B_next": [ + 100, + 100, + 100, + 100, + 100, + 100, + 100, + 100 + ], + "A_cont": [ + [ + 100, + 45518, + 107, + 101, + 221466, + 236764, + 7107, + 506, + 47617, + 33005, + 600, + 2583, + 12175, + 1061, + 13629, + 134480, + 532, + 7595, + 236761, + 1174, + 11639, + 14994, + 22094, + 25671, + 699, + 496, + 73581, + 10946, + 1131, + 496, + 18997, + 50353 + ], + [ + 100, + 45518, + 107, + 101, + 2003, + 506, + 1548, + 529, + 6675, + 6984, + 236764, + 532, + 10769, + 506, + 9113, + 40442, + 12732, + 2342, + 506, + 14628, + 236789, + 236751, + 11519, + 3736, + 11649, + 236761, + 1637, + 506, + 1354, + 16071, + 3426, + 506 + ], + [ + 100, + 45518, + 107, + 101, + 2094, + 8881, + 529, + 4355, + 532, + 136922, + 14004, + 496, + 158490, + 10664, + 236764, + 496, + 158490, + 1972, + 36242, + 684, + 506, + 20133, + 529, + 506, + 7764, + 4319, + 1082, + 506, + 8881, + 529, + 506, + 3768 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 44363, + 497, + 236764, + 25988, + 79639, + 532, + 15127, + 600, + 1093, + 7394, + 75052, + 506, + 72276, + 236761, + 1174, + 11646, + 563, + 506, + 1354, + 529, + 496, + 174801, + 4191, + 236787, + 506, + 86864, + 4920 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 20624, + 29756, + 4319, + 1082, + 16333, + 53350, + 6775, + 506, + 5905, + 21183, + 236761, + 1174, + 20284, + 236764, + 3187, + 11081, + 531, + 618, + 1646, + 623, + 20624, + 236772, + 10633, + 2098, + 2820, + 600, + 506 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 4250, + 15995, + 236761, + 7714, + 123892, + 14736, + 618, + 496, + 24782, + 10021, + 531, + 496, + 21222, + 29280, + 236764, + 496, + 40605, + 10810, + 600, + 33402, + 13947, + 529, + 7635, + 64479, + 496, + 6049, + 21222 + ], + [ + 100, + 45518, + 107, + 101, + 96654, + 1131, + 506, + 3530, + 14787, + 1076, + 13217, + 568, + 25830, + 236768, + 529, + 506, + 3328, + 24990, + 236764, + 18437, + 573, + 18922, + 43313, + 236764, + 15612, + 1757, + 170836, + 236764, + 532, + 9911, + 17046, + 30122 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 23600, + 7002, + 506, + 13258, + 3548, + 529, + 33314, + 532, + 11321, + 11749, + 2342, + 506, + 1440, + 236772, + 6061, + 23811, + 529, + 18226, + 2256, + 236761, + 2195, + 7020, + 506, + 1813, + 4535, + 607, + 496 + ] + ], + "B_cont": [ + [ + 100, + 45518, + 107, + 101, + 221466, + 236764, + 7107, + 506, + 47617, + 33005, + 600, + 2583, + 12175, + 1061, + 13629, + 134480, + 532, + 7595, + 236761, + 1174, + 11639, + 14994, + 22094, + 25671, + 699, + 496, + 73581, + 10946, + 1131, + 496, + 18997, + 50353 + ], + [ + 100, + 45518, + 107, + 101, + 2003, + 506, + 1548, + 529, + 6675, + 6984, + 236764, + 532, + 10769, + 506, + 1354, + 2342, + 506, + 14628, + 236789, + 236751, + 11519, + 16813, + 53350, + 236761, + 1637, + 506, + 40442, + 12732, + 563, + 2708, + 236764, + 506 + ], + [ + 100, + 45518, + 107, + 101, + 2094, + 8881, + 529, + 4355, + 532, + 136922, + 14004, + 496, + 158490, + 10664, + 236764, + 496, + 158490, + 1972, + 36242, + 684, + 506, + 20133, + 529, + 506, + 7764, + 4319, + 1082, + 506, + 8881, + 529, + 506, + 3768 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 44363, + 497, + 236764, + 25988, + 79639, + 532, + 15127, + 600, + 1093, + 7394, + 75052, + 506, + 72276, + 236761, + 1174, + 11646, + 563, + 506, + 1354, + 529, + 496, + 174801, + 4191, + 236787, + 506, + 86864, + 4920 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 20624, + 29756, + 4319, + 1082, + 16333, + 53350, + 6775, + 506, + 5905, + 21183, + 236761, + 1174, + 20284, + 236764, + 3187, + 11081, + 531, + 618, + 1646, + 623, + 20624, + 236772, + 10633, + 2098, + 2820, + 600, + 506 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 4250, + 15995, + 236761, + 7714, + 123892, + 14736, + 618, + 496, + 24782, + 10021, + 531, + 496, + 21222, + 29280, + 236764, + 496, + 40605, + 10810, + 36486, + 13947, + 529, + 7635, + 684, + 496, + 21222, + 8858, + 529 + ], + [ + 100, + 45518, + 107, + 101, + 96654, + 1131, + 506, + 3530, + 14787, + 1076, + 13217, + 568, + 25830, + 236768, + 529, + 506, + 3328, + 24990, + 236764, + 18437, + 573, + 18922, + 43313, + 236764, + 15612, + 1757, + 170836, + 236764, + 532, + 9911, + 17046, + 30122 + ], + [ + 100, + 45518, + 107, + 101, + 1390, + 23600, + 7002, + 506, + 13258, + 3548, + 529, + 33314, + 532, + 11321, + 11749, + 2342, + 506, + 1440, + 236772, + 6061, + 23811, + 529, + 18226, + 2256, + 236761, + 2195, + 7020, + 506, + 1813, + 4535, + 607, + 496 + ] + ], + "A_text": [ + "thought\ndioxide, creating the microscopic bubbles that give bread its characteristic crumb and texture. This scientific revolution transformed baking from a mystical craft into a precise culinary", + "thought\nby the number of operations performed, and compare the resulting arithmetic intensity against the hardware's peak performance limits. If the result falls below the", + "thought\nThis cycle of death and rebirth creates a nomadic existence, a nomadic life governed by the chemistry of the earth rather than the cycle of the sun", + "thought\n...wilder, competing molds and bacteria that would otherwise spoil the loaf. This stability is the result of a symbiotic relationship: the lactic acid", + "thought\n...memory bandwidth rather than compute throughput becomes the primary constraint. This phenomenon, often referred to as being \"memory-bound,\" means that the", + "thought\n...bedrock. Each erratic serves as a silent witness to a frozen epoch, a displaced stone that traveled hundreds of miles atop a moving frozen", + "thought\ninstructions into the specific instruction set architecture (ISA) of the target processor, accounting for pipeline depths, cache line alignments, and branch prediction behaviors", + "thought\n...must balance the immediate needs of irrigation and urban consumption against the long-term necessity of flood control. They monitor the water levels with a" + ] +} \ No newline at end of file diff --git a/temp/gemma4_bprefill/artifacts/gate_fluent_r4.json b/temp/gemma4_bprefill/artifacts/gate_fluent_r4.json new file mode 100644 index 0000000..db67621 --- /dev/null +++ b/temp/gemma4_bprefill/artifacts/gate_fluent_r4.json @@ -0,0 +1,35 @@ +{ + "first_match": [ + true, + true, + true, + true, + true, + true, + true, + true + ], + "cont_match": [ + false, + false, + false, + false, + true, + false, + true, + true + ], + "same_prompt_ok": true, + "perm_next_ok": true, + "perm_cont_ok": true, + "A_text": [ + ",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,", + " measured bandwidth of the..\n.\n.\n.\n.\n.\n.\n.\n.\n..\n..\n..", + ",1,,,,,,,,,,,,.1.1.1,.1.1...1..1", + ". The microbial. The microbial. The microbial. The microbial.. The microbial............1111", + ".\n\n.\n.\n\n.1.1.1.1.1.1.1.1.1.1.1.1.", + ". Today laser altimetry. Today.1.1.1.1.1.1.1.1.1.1.1111", + " to. The final emission stage. The final. The. The...................", + ". Reservoir managers, city planners, and insurance. Reservoir managers,, city planners, and. Reservoir managers,,,,,,11111" + ] +} \ No newline at end of file diff --git a/temp/gemma4_bprefill/artifacts/gates_12b_fixed.json b/temp/gemma4_bprefill/artifacts/gates_12b_fixed.json new file mode 100644 index 0000000..046acd1 --- /dev/null +++ b/temp/gemma4_bprefill/artifacts/gates_12b_fixed.json @@ -0,0 +1,675 @@ +{ + "model": "gemma4-12b-unified", + "batch": 8, + "window": 1024, + "T": 128, + "chunk": 64, + "gate1": { + "first_match": [ + true, + true, + true, + true, + true, + true, + true, + true + ], + "cont_match": [ + false, + false, + false, + true, + false, + false, + false, + true + ], + "chunk_vs_single": [ + true, + true, + true, + true, + true, + true, + true, + true + ], + "tokens_in_vocab": true, + "base_next": [ + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761 + ], + "new_next": [ + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761 + ], + "base_cont": [ + [ + 236761, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107 + ], + [ + 236761, + 107, + 1, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107 + ], + [ + 236761, + 107, + 1, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107, + 236865, + 107 + ], + [ + 236761, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770 + ], + [ + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ], + [ + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ], + [ + 236761, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761 + ], + [ + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ] + ], + "new_cont": [ + [ + 236761, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107, + 236772, + 107 + ], + [ + 236761, + 1, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ], + [ + 236761, + 107, + 1, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997, + 107, + 997 + ], + [ + 236761, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770, + 236761, + 107, + 236770 + ], + [ + 236761, + 107, + 1, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761 + ], + [ + 236761, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107, + 107 + ], + [ + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761 + ], + [ + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ] + ] + }, + "gate5": { + "dec_a_tok_s": 10.918478759933436, + "dec_b_tok_s": 10.918008096310801, + "ratio": 0.9999568929305096 + }, + "gate1b_same_prompt_next": [ + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761 + ], + "ttft": [ + { + "T": 128, + "chunk": 0, + "A_ms": 93360.30879200007, + "B_ms": 18275.000125000133, + "speedup": 5.108635193073598, + "A_all": [ + 93350.04274999995, + 93387.93499999997, + 93311.05308300005, + 93348.08662499995, + 93360.30879200007, + 93415.12629099998, + 93416.04237499996 + ], + "B_all": [ + 18277.436875000036, + 18256.26529200008, + 18372.523916999966, + 18446.62645799997, + 18275.000125000133, + 18251.546125000004, + 18231.25062500003 + ] + }, + { + "T": 128, + "chunk": 64, + "A_ms": 93360.30879200007, + "B_ms": 16856.863290999856, + "speedup": 5.538415254387607, + "A_all": [ + 93350.04274999995, + 93387.93499999997, + 93311.05308300005, + 93348.08662499995, + 93360.30879200007, + 93415.12629099998, + 93416.04237499996 + ], + "B_all": [ + 16855.45275000004, + 16941.855165999983, + 16849.549499999965, + 16850.130166999974, + 16856.863290999856, + 16859.274708000157, + 16894.235708999986 + ] + } + ] +} \ No newline at end of file diff --git a/temp/gemma4_bprefill/artifacts/seq_paths_chat.json b/temp/gemma4_bprefill/artifacts/seq_paths_chat.json new file mode 100644 index 0000000..1558e4e --- /dev/null +++ b/temp/gemma4_bprefill/artifacts/seq_paths_chat.json @@ -0,0 +1,586 @@ +[ + { + "prompt": 0, + "match": true, + "div_idx": -1, + "c1": [ + 100, + 45518, + 107, + 101, + 221466, + 236764, + 7107, + 506, + 47617, + 33005, + 600, + 2583, + 12175, + 1061, + 13629, + 134480, + 532, + 7595, + 236761, + 1174, + 11639, + 14994, + 22094, + 25671, + 699, + 496, + 73581, + 10946, + 1131, + 496, + 18997, + 50353 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 221466, + 236764, + 7107, + 506, + 47617, + 33005, + 600, + 2583, + 12175, + 1061, + 13629, + 134480, + 532, + 7595, + 236761, + 1174, + 11639, + 14994, + 22094, + 25671, + 699, + 496, + 73581, + 10946, + 1131, + 496, + 18997, + 50353 + ] + }, + { + "prompt": 1, + "match": false, + "div_idx": 14, + "c1": [ + 100, + 45518, + 107, + 101, + 2003, + 506, + 1548, + 529, + 6675, + 6984, + 236764, + 532, + 10769, + 506, + 9113, + 40442, + 12732, + 2342, + 506, + 14628, + 236789, + 236751, + 11519, + 3736, + 11649, + 236761, + 1637, + 506, + 1354, + 16071, + 3426, + 506 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 2003, + 506, + 1548, + 529, + 6675, + 6984, + 236764, + 532, + 10769, + 506, + 1354, + 2342, + 506, + 14628, + 236789, + 236751, + 11519, + 16813, + 53350, + 236761, + 1637, + 506, + 40442, + 12732, + 563, + 2708, + 236764, + 506 + ] + }, + { + "prompt": 2, + "match": true, + "div_idx": -1, + "c1": [ + 100, + 45518, + 107, + 101, + 2094, + 8881, + 529, + 4355, + 532, + 136922, + 14004, + 496, + 158490, + 10664, + 236764, + 496, + 158490, + 1972, + 36242, + 684, + 506, + 20133, + 529, + 506, + 7764, + 4319, + 1082, + 506, + 8881, + 529, + 506, + 3768 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 2094, + 8881, + 529, + 4355, + 532, + 136922, + 14004, + 496, + 158490, + 10664, + 236764, + 496, + 158490, + 1972, + 36242, + 684, + 506, + 20133, + 529, + 506, + 7764, + 4319, + 1082, + 506, + 8881, + 529, + 506, + 3768 + ] + }, + { + "prompt": 3, + "match": true, + "div_idx": -1, + "c1": [ + 100, + 45518, + 107, + 101, + 1390, + 44363, + 497, + 236764, + 25988, + 79639, + 532, + 15127, + 600, + 1093, + 7394, + 75052, + 506, + 72276, + 236761, + 1174, + 11646, + 563, + 506, + 1354, + 529, + 496, + 174801, + 4191, + 236787, + 506, + 86864, + 4920 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 1390, + 44363, + 497, + 236764, + 25988, + 79639, + 532, + 15127, + 600, + 1093, + 7394, + 75052, + 506, + 72276, + 236761, + 1174, + 11646, + 563, + 506, + 1354, + 529, + 496, + 174801, + 4191, + 236787, + 506, + 86864, + 4920 + ] + }, + { + "prompt": 4, + "match": true, + "div_idx": -1, + "c1": [ + 100, + 45518, + 107, + 101, + 1390, + 20624, + 29756, + 4319, + 1082, + 16333, + 53350, + 6775, + 506, + 5905, + 21183, + 236761, + 1174, + 20284, + 236764, + 3187, + 11081, + 531, + 618, + 1646, + 623, + 20624, + 236772, + 10633, + 2098, + 2820, + 600, + 506 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 1390, + 20624, + 29756, + 4319, + 1082, + 16333, + 53350, + 6775, + 506, + 5905, + 21183, + 236761, + 1174, + 20284, + 236764, + 3187, + 11081, + 531, + 618, + 1646, + 623, + 20624, + 236772, + 10633, + 2098, + 2820, + 600, + 506 + ] + }, + { + "prompt": 5, + "match": true, + "div_idx": -1, + "c1": [ + 100, + 45518, + 107, + 101, + 1390, + 4250, + 15995, + 236761, + 7714, + 123892, + 14736, + 618, + 496, + 24782, + 10021, + 531, + 496, + 21222, + 29280, + 236764, + 496, + 40605, + 10810, + 36486, + 13947, + 529, + 7635, + 684, + 496, + 21222, + 8858, + 529 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 1390, + 4250, + 15995, + 236761, + 7714, + 123892, + 14736, + 618, + 496, + 24782, + 10021, + 531, + 496, + 21222, + 29280, + 236764, + 496, + 40605, + 10810, + 36486, + 13947, + 529, + 7635, + 684, + 496, + 21222, + 8858, + 529 + ] + }, + { + "prompt": 6, + "match": true, + "div_idx": -1, + "c1": [ + 100, + 45518, + 107, + 101, + 96654, + 1131, + 506, + 3530, + 14787, + 1076, + 13217, + 568, + 25830, + 236768, + 529, + 506, + 3328, + 24990, + 236764, + 18437, + 573, + 18922, + 43313, + 236764, + 15612, + 1757, + 170836, + 236764, + 532, + 9911, + 17046, + 30122 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 96654, + 1131, + 506, + 3530, + 14787, + 1076, + 13217, + 568, + 25830, + 236768, + 529, + 506, + 3328, + 24990, + 236764, + 18437, + 573, + 18922, + 43313, + 236764, + 15612, + 1757, + 170836, + 236764, + 532, + 9911, + 17046, + 30122 + ] + }, + { + "prompt": 7, + "match": false, + "div_idx": 4, + "c1": [ + 100, + 45518, + 107, + 101, + 1390, + 23600, + 7002, + 506, + 13258, + 3548, + 529, + 33314, + 532, + 11321, + 11749, + 2342, + 506, + 1440, + 236772, + 6061, + 23811, + 529, + 18226, + 2256, + 236761, + 2195, + 7020, + 506, + 1813, + 4535, + 607, + 496 + ], + "c2": [ + 100, + 45518, + 107, + 101, + 14625, + 7020, + 506, + 1813, + 4535, + 529, + 12566, + 56011, + 236764, + 1921, + 7002, + 506, + 13258, + 3548, + 529, + 33314, + 532, + 11321, + 11749, + 2342, + 506, + 1440, + 236772, + 6061, + 23811, + 529, + 18226, + 2256 + ] + } +] \ No newline at end of file diff --git a/temp/gemma4_bprefill/artifacts/seq_paths_fluent.json b/temp/gemma4_bprefill/artifacts/seq_paths_fluent.json new file mode 100644 index 0000000..092dc90 --- /dev/null +++ b/temp/gemma4_bprefill/artifacts/seq_paths_fluent.json @@ -0,0 +1,586 @@ +[ + { + "prompt": 0, + "match": false, + "div_idx": 9, + "c1": [ + 236764, + 236764, + 236764, + 236764, + 3095, + 531, + 531, + 531, + 531, + 210285, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723, + 219723 + ], + "c2": [ + 236764, + 236764, + 236764, + 236764, + 3095, + 531, + 531, + 531, + 531, + 199, + 236764, + 236764, + 236764, + 236764, + 581, + 581, + 581, + 581, + 581, + 581, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771, + 236771 + ] + }, + { + "prompt": 1, + "match": true, + "div_idx": -1, + "c1": [ + 8434, + 29756, + 529, + 506, + 236761, + 255999, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ], + "c2": [ + 8434, + 29756, + 529, + 506, + 236761, + 255999, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107, + 236761, + 107 + ] + }, + { + "prompt": 2, + "match": false, + "div_idx": 1, + "c1": [ + 236764, + 6272, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770 + ], + "c2": [ + 236764, + 236770, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236764, + 236761, + 236770, + 236761, + 236761, + 236761, + 236770, + 236761, + 236761, + 236761, + 236770, + 236761 + ] + }, + { + "prompt": 3, + "match": false, + "div_idx": 28, + "c1": [ + 236761, + 669, + 47682, + 236761, + 669, + 47682, + 236761, + 669, + 47682, + 236761, + 669, + 47682, + 236761, + 236761, + 669, + 47682, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236770, + 236770, + 236770, + 236770 + ], + "c2": [ + 236761, + 669, + 47682, + 236761, + 669, + 47682, + 236761, + 669, + 47682, + 236761, + 669, + 47682, + 236761, + 236761, + 669, + 47682, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236770 + ] + }, + { + "prompt": 4, + "match": true, + "div_idx": -1, + "c1": [ + 236761, + 108, + 236761, + 107, + 236761, + 107, + 107, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761 + ], + "c2": [ + 236761, + 108, + 236761, + 107, + 236761, + 107, + 107, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761, + 236770, + 236761 + ] + }, + { + "prompt": 5, + "match": false, + "div_idx": 0, + "c1": [ + 236761, + 10950, + 13678, + 4466, + 74917, + 236761, + 10950, + 236761, + 108, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770 + ], + "c2": [ + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770, + 236770 + ] + }, + { + "prompt": 6, + "match": true, + "div_idx": -1, + "c1": [ + 531, + 236761, + 669, + 1626, + 16364, + 5552, + 236761, + 669, + 1626, + 236761, + 669, + 236761, + 669, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761 + ], + "c2": [ + 531, + 236761, + 669, + 1626, + 16364, + 5552, + 236761, + 669, + 1626, + 236761, + 669, + 236761, + 669, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761 + ] + }, + { + "prompt": 7, + "match": true, + "div_idx": -1, + "c1": [ + 236761, + 80177, + 16810, + 236764, + 3207, + 67172, + 236764, + 532, + 8657, + 236761, + 80177, + 16810, + 236764, + 236764, + 3207, + 67172, + 236764, + 532, + 236761, + 80177, + 16810, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236770, + 236770, + 236770, + 236770, + 236770 + ], + "c2": [ + 236761, + 80177, + 16810, + 236764, + 3207, + 67172, + 236764, + 532, + 8657, + 236761, + 80177, + 16810, + 236764, + 236764, + 3207, + 67172, + 236764, + 532, + 236761, + 80177, + 16810, + 236764, + 236764, + 236764, + 236764, + 236764, + 236764, + 236770, + 236770, + 236770, + 236770, + 236770 + ] + } +] \ No newline at end of file diff --git a/temp/gemma4_bprefill/artifacts/ss_new_r4.json b/temp/gemma4_bprefill/artifacts/ss_new_r4.json new file mode 100644 index 0000000..f2220e2 --- /dev/null +++ b/temp/gemma4_bprefill/artifacts/ss_new_r4.json @@ -0,0 +1,73 @@ +{ + "dylib": "/Users/derek/sk-gemma-bprefill-r4/build/libsk.dylib", + "coherence_ids": [ + 100, + 45518, + 107, + 101, + 236776, + 85618, + 529, + 18763, + 236764, + 1133, + 22303, + 7613, + 236764, + 107, + 902, + 496, + 7116, + 13392, + 1298, + 506, + 37676, + 2727, + 236761, + 107, + 236776, + 181283, + 529, + 2173, + 580, + 496, + 3761, + 529, + 11261, + 236764, + 107, + 60776, + 573, + 1813, + 531, + 18259, + 699, + 1061, + 6927, + 236761, + 108, + 11407, + 3952, + 506 + ], + "coherence_text": "thought\nA mound of flour, like fallen snow,\nIn a wooden bowl where the shadows flow.\nA dusting of white on a surface of grain,\nWaiting for water to wake from its rain.\n\nThen comes the", + "decode_ids": [ + 236770, + 236770, + 236770, + 236761, + 236770, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236761, + 236770, + 236770, + 236770, + 236761, + 236761 + ] +} \ No newline at end of file diff --git a/temp/gemma4_bprefill/gate_chat.py b/temp/gemma4_bprefill/gate_chat.py new file mode 100644 index 0000000..b090aac --- /dev/null +++ b/temp/gemma4_bprefill/gate_chat.py @@ -0,0 +1,107 @@ +"""Gate 2b: chat-templated FLUENT prompts (sharp argmax — the it-model's regime). + +Same A/B as gate_fluent (A = token-by-token lockstep, B = sk_gemma4_prefill_batched) +but each lane's prompt is a full gemma chat turn ("continue this passage") spliced +to EXACTLY T tokens with the model-turn cue intact, so greedy continuations are +fluent prose, not repetition loops with knife-edge argmax. +""" +import os, sys, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r4"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--batch", type=int, default=8) +ap.add_argument("--T", type=int, default=128) +ap.add_argument("--chunk", type=int, default=64) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +m = sk.load("gemma4-12b-unified", batch=args.batch, seq_max=128, cache_max=512, + window=512) +print(f"[setup] batch={args.batch} window={m.cfg.window}", flush=True) + +from prompts_fluent import TEXTS + +def chat_ids(text): + return list(m.tokenizer.chat( + [{"role": "user", "content": "Continue this passage:\n\n" + text}], bos=True)) + +# Common suffix length S of the template (model-turn cue) — splice point. +a, b = chat_ids(TEXTS[0]), chat_ids(TEXTS[1]) +S = 0 +while S < min(len(a), len(b)) and a[-1 - S] == b[-1 - S]: + S += 1 +print(f"[setup] template suffix tokens S={S}", flush=True) + +rows = [] +for t in TEXTS: + full = chat_ids(t) + assert len(full) >= args.T, f"templated prompt too short: {len(full)}" + rows.append(full[: args.T - S] + full[-S:]) +ids = np.ascontiguousarray(np.array(rows, dtype=np.int32)) + +def prefill_tbt(p): + out = None + for t in range(p.shape[1]): + out = m.forward_batched(np.ascontiguousarray(p[:, t])) + return out + +def decode_n(first, n): + toks = [np.array(first, dtype=np.int32).copy()] + for _ in range(n - 1): + toks.append(m.forward_batched(toks[-1]).astype(np.int32).copy()) + return np.stack(toks, axis=1) + +m.reset() +base_next = prefill_tbt(ids) +base_cont = decode_n(base_next, args.cont) +print("[chat] A done", flush=True) + +m.reset() +new_next = m.prefill_batched(ids, chunk_size=args.chunk) +new_cont = decode_n(new_next, args.cont) +print("[chat] B done", flush=True) + +first_match = [int(base_next[b_]) == int(new_next[b_]) for b_ in range(args.batch)] +lane_match = [bool((base_cont[b_] == new_cont[b_]).all()) for b_ in range(args.batch)] +print(f"[chat] first-token match per lane: {first_match}") +print(f"[chat] {args.cont}-token continuation match per lane: {lane_match}", flush=True) +for b_ in range(args.batch): + tag = "OK " if lane_match[b_] else "DIV" + print(f"[chat] lane {b_} {tag} A text: {m.tokenizer.decode(base_cont[b_].tolist())!r}") + if not lane_match[b_]: + x, y = base_cont[b_].tolist(), new_cont[b_].tolist() + d = next(i for i in range(len(x)) if x[i] != y[i]) + print(f"[chat] lane {b_} diverges at idx {d}: A={x[max(0,d-2):d+3]} B={y[max(0,d-2):d+3]}") + print(f"[chat] lane {b_} B text: {m.tokenizer.decode(new_cont[b_].tolist())!r}") + +# identical-prompt + permutation invariants on the templated prompts. +ids_same = np.tile(ids[0], (args.batch, 1)) +m.reset() +same_next = m.prefill_batched(ids_same, chunk_size=args.chunk) +same_ok = bool((same_next == same_next[0]).all()) +print(f"[chat] identical prompts -> identical next tokens: {same_ok}", flush=True) + +perm = np.array([3, 1, 4, 0, 7, 5, 2, 6]) +m.reset() +pnext = m.prefill_batched(np.ascontiguousarray(ids[perm]), chunk_size=args.chunk) +pcont = decode_n(pnext, 8) +m.reset() +qnext = m.prefill_batched(ids, chunk_size=args.chunk) +qcont = decode_n(qnext, 8) +perm_next_ok = bool((np.array(pnext) == np.array(qnext)[perm]).all()) +perm_cont_ok = bool((pcont == qcont[perm]).all()) +print(f"[perm] next tokens permute exactly: {perm_next_ok}") +print(f"[perm] 8-tok continuations permute exactly: {perm_cont_ok}", flush=True) + +res = dict(first_match=first_match, cont_match=lane_match, same_prompt_ok=same_ok, + perm_next_ok=perm_next_ok, perm_cont_ok=perm_cont_ok, + A_next=[int(x) for x in base_next], B_next=[int(x) for x in new_next], + A_cont=base_cont.tolist(), B_cont=new_cont.tolist(), + A_text=[m.tokenizer.decode(base_cont[b_].tolist()) for b_ in range(args.batch)]) +if args.json_out: + json.dump(res, open(args.json_out, "w"), indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/gate_fluent.py b/temp/gemma4_bprefill/gate_fluent.py new file mode 100644 index 0000000..8ba03d3 --- /dev/null +++ b/temp/gemma4_bprefill/gate_fluent.py @@ -0,0 +1,106 @@ +"""Gate 1 with FLUENT prompts: 8 unique long passages (no tiling), BOS at pos 0. + +Checks (per WIN gate): + - 8/8 per-lane exact match (first token + 32-tok continuation) B vs A, + where A = token-by-token lockstep (M=8 per step, never touches the + BM=32-vs-/64 row-grid sites) and B = sk_gemma4_prefill_batched. + - identical-prompt invariant, lane-permutation invariant. + - prints A-side decoded text per lane so fluency is verifiable. +""" +import os, sys, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r3"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--batch", type=int, default=8) +ap.add_argument("--T", type=int, default=128) +ap.add_argument("--chunk", type=int, default=64) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +m = sk.load("gemma4-12b-unified", batch=args.batch, seq_max=128, cache_max=512, + window=512) +print(f"[setup] batch={args.batch} window={m.cfg.window}", flush=True) + +TEXTS = [ + "The history of bread baking stretches back over ten thousand years, beginning with flat unleavened cakes cooked on hot stones beside open fires. Ancient Egyptian bakers discovered that dough left to rest would rise on its own, captured wild yeasts turning a dense paste into an airy loaf. Roman bakeries industrialized the craft with large masonry ovens and professional guilds, and by the Middle Ages nearly every European village supported a communal oven where families brought their shaped loaves each morning. The chemistry behind all of this remained mysterious until the nineteenth century, when scientists finally described how yeast ferments sugars into carbon dioxide, stretching gluten networks that bakers had been kneading by intuition for millennia. Modern artisans now combine that scientific understanding with", + "Metal compute shaders allow a programmer to dispatch thousands of threadgroups across the GPU, each cooperating through fast threadgroup memory while reading and writing device buffers. A well designed kernel keeps its arithmetic intensity high enough to hide memory latency, streaming tiles of data through shared storage so that every byte fetched from DRAM is reused many times before being discarded. On Apple Silicon the unified memory architecture removes explicit copies between processor and accelerator, but the bandwidth ceiling still dominates performance for most inference workloads. Profiling therefore begins with a simple roofline analysis: count the bytes each kernel must move, divide by the measured bandwidth of the chip, and compare that bound against the observed execution time to decide whether", + "In the deep ocean, hydrothermal vents support entire ecosystems that never see sunlight, powered instead by chemosynthetic bacteria that oxidize hydrogen sulfide gushing from the seafloor. Giant tube worms cluster around the vents in dense thickets, lacking mouths and digestive tracts entirely, nourished by symbiotic microbes housed within their tissues. Crabs, shrimp, and pale octopuses patrol the mineral chimneys, grazing on bacterial mats that coat every surface. When a vent finally goes extinct, the community collapses within months, yet larvae carried on deep currents somehow locate new vents tens of kilometers away and rebuild the entire assemblage. Biologists studying these habitats argue that similar chemical gardens beneath the ice of distant moons could plausibly", + "A well tuned sourdough starter doubles in volume within four to six hours of feeding, producing a pleasant aroma of ripe fruit and yogurt rather than the sharp smell of acetone that signals neglect. Maintaining that vigor requires a steady rhythm: discard most of the culture, refresh it with equal weights of flour and water, and hold it at a temperature where the yeasts and lactic acid bacteria stay in balance. Bakers who keep their starter in the refrigerator slow this cycle to a weekly feeding, trading some liveliness for convenience. The microbial community inside a mature starter is remarkably stable, resisting invasion by stray organisms because its acidity and competitive ecology leave no niche unfilled, which explains why", + "The transformer architecture replaced recurrence with attention, letting every token attend directly to every other token in the sequence instead of squeezing history through a fixed size hidden state. This change unlocked massive parallelism during training, since whole sequences could be processed at once on modern accelerators rather than step by step. The cost is quadratic scaling in sequence length, which has driven a decade of research into sparse patterns, sliding windows, linear approximations, and key value caching strategies. During inference the bottleneck shifts: generating each new token requires reading the entire stack of cached keys and values, so memory bandwidth rather than raw arithmetic throughput usually determines how quickly a large language model can", + "Glaciers carve valleys over millennia, grinding bedrock into fine flour that turns meltwater lakes a striking shade of turquoise blue. As the ice advances it plucks boulders from the valley walls and drags them along its base, scouring deep U shaped troughs that remain long after the climate warms. Moraines of tumbled rock mark each pause in the retreat, and stranded blocks of ice leave kettle ponds dotting the outwash plain. Geologists read these landforms like a written record, reconstructing the extent of ancient ice sheets from ridgelines and erratic boulders perched far from their parent outcrops. Today laser altimetry and satellite gravimetry extend that record forward, measuring yearly losses of ice mass that", + "Compilers translate high level source code into machine instructions through stages of parsing, optimization, and code generation, each built on decades of formal theory. The front end checks syntax and types, lowering the program into an intermediate representation that captures its meaning while discarding surface details. Optimization passes then rewrite this representation, folding constants, eliminating dead branches, hoisting loop invariant work, and vectorizing inner loops where the hardware allows. Register allocation maps an unbounded supply of virtual values onto a handful of physical registers, spilling the overflow to the stack as cheaply as possible. The final emission stage schedules instructions to keep the processor pipelines full, because even a perfectly optimized sequence of operations will", + "The annual monsoon arrives on the southwest coast in early June, bringing weeks of heavy rain that replenish rivers and aquifers after the long dry season. Farmers time their sowing to the first reliable downpours, and an early or late onset can shift harvests across an entire subcontinent. The system itself is a vast heat engine: land warms faster than ocean in spring, drawing moist air inland where it rises over the mountains and releases its water. Forecasters track sea surface temperatures thousands of kilometers away because shifts in the Pacific can strengthen or starve the circulation months in advance. Reservoir managers, city planners, and insurance companies all build their yearly calendars around the", +] + +rows = [] +for t in TEXTS: + ids = m.tokenizer.encode(t) # adds BOS + assert len(ids) >= args.T, f"prompt too short: {len(ids)}" + rows.append(ids[:args.T]) +ids = np.ascontiguousarray(np.array(rows, dtype=np.int32)) + +def prefill_tbt(p): + out = None + for t in range(p.shape[1]): + out = m.forward_batched(np.ascontiguousarray(p[:, t])) + return out + +def decode_n(first, n): + toks = [np.array(first, dtype=np.int32).copy()] + for _ in range(n - 1): + toks.append(m.forward_batched(toks[-1]).astype(np.int32).copy()) + return np.stack(toks, axis=1) + +m.reset() +base_next = prefill_tbt(ids) +base_cont = decode_n(base_next, args.cont) +print("[fluent] A done", flush=True) + +m.reset() +new_next = m.prefill_batched(ids, chunk_size=args.chunk) +new_cont = decode_n(new_next, args.cont) +print("[fluent] B done", flush=True) + +first_match = [int(base_next[b]) == int(new_next[b]) for b in range(args.batch)] +lane_match = [bool((base_cont[b] == new_cont[b]).all()) for b in range(args.batch)] +print(f"[fluent] first-token match per lane: {first_match}") +print(f"[fluent] {args.cont}-token continuation match per lane: {lane_match}", flush=True) +for b in range(args.batch): + tag = "OK " if lane_match[b] else "DIV" + print(f"[fluent] lane {b} {tag} A text: {m.tokenizer.decode(base_cont[b].tolist())!r}") + if not lane_match[b]: + a, c = base_cont[b].tolist(), new_cont[b].tolist() + d = next(i for i in range(len(a)) if a[i] != c[i]) + print(f"[fluent] lane {b} diverges at idx {d}: A={a[max(0,d-2):d+3]} B={c[max(0,d-2):d+3]}") + print(f"[fluent] lane {b} B text: {m.tokenizer.decode(new_cont[b].tolist())!r}") + +# identical-prompt invariant +ids_same = np.tile(ids[0], (args.batch, 1)) +m.reset() +same_next = m.prefill_batched(ids_same, chunk_size=args.chunk) +same_ok = bool((same_next == same_next[0]).all()) +print(f"[fluent] identical prompts -> identical next tokens: {same_ok}", flush=True) + +# lane-permutation invariant +perm = np.array([3, 1, 4, 0, 7, 5, 2, 6]) +m.reset() +pnext = m.prefill_batched(np.ascontiguousarray(ids[perm]), chunk_size=args.chunk) +pcont = decode_n(pnext, 8) +m.reset() +qnext = m.prefill_batched(ids, chunk_size=args.chunk) +qcont = decode_n(qnext, 8) +perm_next_ok = bool((np.array(pnext) == np.array(qnext)[perm]).all()) +perm_cont_ok = bool((pcont == qcont[perm]).all()) +print(f"[perm] next tokens permute exactly: {perm_next_ok}") +print(f"[perm] 8-tok continuations permute exactly: {perm_cont_ok}", flush=True) + +res = dict(first_match=first_match, cont_match=lane_match, same_prompt_ok=same_ok, + perm_next_ok=perm_next_ok, perm_cont_ok=perm_cont_ok, + A_text=[m.tokenizer.decode(base_cont[b].tolist()) for b in range(args.batch)]) +if args.json_out: + json.dump(res, open(args.json_out, "w"), indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/gate_real_prompts.py b/temp/gemma4_bprefill/gate_real_prompts.py new file mode 100644 index 0000000..f077f67 --- /dev/null +++ b/temp/gemma4_bprefill/gate_real_prompts.py @@ -0,0 +1,99 @@ +"""Gate 1 with REAL text prompts (peaked logit distributions) + bitwise +lane-permutation test. + +Random-token prompts give degenerate newline-heavy continuations where the +pre-existing seq1-vs-seqN attention reduction-order noise flips argmax ties; +real prompts are the representative serving case. +""" +import os, sys, time, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r2"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--batch", type=int, default=8) +ap.add_argument("--T", type=int, default=128) +ap.add_argument("--chunk", type=int, default=64) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +m = sk.load("gemma4-12b-unified", batch=args.batch, seq_max=128, cache_max=512, + window=512) +print(f"[setup] batch={args.batch} window={m.cfg.window}", flush=True) + +TEXTS = [ + "The history of bread baking stretches back over ten thousand years, beginning with flat unleavened cakes cooked on hot stones. ", + "Metal compute shaders allow a programmer to dispatch thousands of threadgroups across the GPU, each cooperating through threadgroup memory. ", + "In the deep ocean, hydrothermal vents support entire ecosystems that never see sunlight, powered instead by chemosynthetic bacteria. ", + "A well-tuned sourdough starter doubles in volume within four to six hours of feeding, producing a pleasant aroma of ripe fruit and yogurt. ", + "The transformer architecture replaced recurrence with attention, letting every token attend directly to every other token in the sequence. ", + "Glaciers carve valleys over millennia, grinding bedrock into fine flour that turns meltwater lakes a striking shade of turquoise blue. ", + "Compilers translate high level source code into machine instructions through stages of parsing, optimization, and code generation. ", + "The annual monsoon arrives on the southwest coast in early June, bringing weeks of heavy rain that replenish rivers and aquifers. ", +] + +rows = [] +for t in TEXTS: + ids = m.tokenizer.encode(t * 10) + assert len(ids) >= args.T, f"prompt too short: {len(ids)}" + rows.append(ids[:args.T]) +ids = np.ascontiguousarray(np.array(rows, dtype=np.int32)) + +def prefill_tbt(p): + out = None + for t in range(p.shape[1]): + out = m.forward_batched(np.ascontiguousarray(p[:, t])) + return out + +def decode_n(first, n): + toks = [np.array(first, dtype=np.int32).copy()] + for _ in range(n - 1): + toks.append(m.forward_batched(toks[-1]).astype(np.int32).copy()) + return np.stack(toks, axis=1) + +m.reset() +base_next = prefill_tbt(ids) +base_cont = decode_n(base_next, args.cont) +print("[real] A done", flush=True) + +m.reset() +new_next = m.prefill_batched(ids, chunk_size=args.chunk) +new_cont = decode_n(new_next, args.cont) +print("[real] B done", flush=True) + +first_match = [int(base_next[b]) == int(new_next[b]) for b in range(args.batch)] +lane_match = [bool((base_cont[b] == new_cont[b]).all()) for b in range(args.batch)] +print(f"[real] first-token match per lane: {first_match}") +print(f"[real] {args.cont}-token continuation match per lane: {lane_match}", flush=True) +for b in range(args.batch): + if not lane_match[b]: + a, c = base_cont[b].tolist(), new_cont[b].tolist() + d = next(i for i in range(len(a)) if a[i] != c[i]) + print(f"[real] lane {b} diverges at idx {d}: A={a[max(0,d-2):d+3]} B={c[max(0,d-2):d+3]}") + print(f"[real] lane {b} A text: {m.tokenizer.decode(base_cont[b].tolist())!r}") + print(f"[real] lane {b} B text: {m.tokenizer.decode(new_cont[b].tolist())!r}") + +# Bitwise lane-permutation test (same code path both runs): permuting the +# prompts must permute next tokens AND short continuations exactly — any +# cross-lane leakage (KV slice / row indexing) breaks this. +perm = np.array([3, 1, 4, 0, 7, 5, 2, 6]) +m.reset() +pnext = m.prefill_batched(np.ascontiguousarray(ids[perm]), chunk_size=args.chunk) +pcont = decode_n(pnext, 8) +m.reset() +qnext = m.prefill_batched(ids, chunk_size=args.chunk) +qcont = decode_n(qnext, 8) +perm_next_ok = bool((np.array(pnext) == np.array(qnext)[perm]).all()) +perm_cont_ok = bool((pcont == qcont[perm]).all()) +print(f"[perm] next tokens permute exactly: {perm_next_ok}") +print(f"[perm] 8-token continuations permute exactly: {perm_cont_ok}", flush=True) + +res = dict(first_match=first_match, cont_match=lane_match, + perm_next_ok=perm_next_ok, perm_cont_ok=perm_cont_ok, + base_next=[int(x) for x in base_next], new_next=[int(x) for x in new_next], + base_cont=base_cont.tolist(), new_cont=new_cont.tolist()) +if args.json_out: + json.dump(res, open(args.json_out, "w"), indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/gates_ttft.py b/temp/gemma4_bprefill/gates_ttft.py new file mode 100644 index 0000000..b950504 --- /dev/null +++ b/temp/gemma4_bprefill/gates_ttft.py @@ -0,0 +1,157 @@ +"""gemma4-12B-unified batched-prefill gates + serving-TTFT A/B (single process, one handle). + +A = baseline serving prefill: token-by-token lockstep (forward_batched seq=1, T steps) +B = new sk_gemma4_prefill_batched (chunked M=batch*seq prefill) + +Gate 1 (lane isolation): per-lane first token + 32-token greedy continuation after B +must match A's, per lane; chunk=64 == single-chunk; identical prompts -> identical lanes. +Gate 4 (TTFT): reset -> all-lanes-first-token, 2 warmups + 7 reps median, 0.3s gaps. +Gate 5 (decode aggregate): 32 lockstep steps after A-prefill vs after B-prefill, +-2%. + +Env: SK_DYLIB, SK_METALLIB=/nonexistent, SK_METAL_SRC_FALLBACK, SK_TREE, + SK_GEMMA4_BODY_Q4K=1, SK_GEMMA4_EMBED_Q8=1. +""" +import os, sys, time, json, argparse, statistics + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r2"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--batch", type=int, default=8) +ap.add_argument("--seq_max", type=int, default=128) +ap.add_argument("--cache_max", type=int, default=512) +ap.add_argument("--window", type=int, default=0) # 0 = model default (1024) +ap.add_argument("--T", type=int, default=128) +ap.add_argument("--chunk", type=int, default=64) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--reps", type=int, default=7) +ap.add_argument("--skip_gates", action="store_true") +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +over = dict(batch=args.batch, seq_max=args.seq_max, cache_max=args.cache_max) +if args.window: + over["window"] = args.window +t0 = time.perf_counter() +m = sk.load("gemma4-12b-unified", **over) +print(f"[setup] loaded in {time.perf_counter()-t0:.1f}s batch={args.batch} " + f"seq_max={args.seq_max} cache_max={args.cache_max} window={m.cfg.window}", flush=True) + +rng = np.random.default_rng(7) +def make_prompts(T): + return np.ascontiguousarray(rng.integers(10, 200000, size=(args.batch, T)).astype(np.int32)) + +def prefill_tbt(ids): + out = None + for t in range(ids.shape[1]): + out = m.forward_batched(np.ascontiguousarray(ids[:, t])) + return out + +def decode_n(first, n): + toks = [np.array(first, dtype=np.int32).copy()] + cur = toks[0] + for _ in range(n - 1): + cur = m.forward_batched(cur).astype(np.int32) + toks.append(cur.copy()) + return np.stack(toks, axis=1) # (batch, n) + +results = {"model": "gemma4-12b-unified", "batch": args.batch, + "window": int(m.cfg.window), "T": args.T, "chunk": args.chunk} + +if not args.skip_gates: + # ---- Gate 1: lane isolation --------------------------------------------- + ids = make_prompts(args.T) + print(f"[gate1] T={args.T} chunk={args.chunk} cont={args.cont}", flush=True) + + m.reset() + t0 = time.perf_counter() + base_next = prefill_tbt(ids) + print(f"[gate1] tbt prefill done in {time.perf_counter()-t0:.1f}s", flush=True) + t0 = time.perf_counter() + base_cont = decode_n(base_next, args.cont) + dt_dec_a = time.perf_counter() - t0 + print(f"[gate1] A continuation done in {dt_dec_a:.1f}s", flush=True) + + m.reset() + t0 = time.perf_counter() + new_next = m.prefill_batched(ids, chunk_size=args.chunk) + print(f"[gate1] batched prefill done in {time.perf_counter()-t0:.1f}s", flush=True) + t0 = time.perf_counter() + new_cont = decode_n(new_next, args.cont) + dt_dec_b = time.perf_counter() - t0 + + m.reset() + new_next_1c = m.prefill_batched(ids, chunk_size=0) # single chunk (<= seq_max) + + lane_match = [bool((base_cont[b] == new_cont[b]).all()) for b in range(args.batch)] + first_match = [int(base_next[b]) == int(new_next[b]) for b in range(args.batch)] + chunk_vs_1c = [int(new_next[b]) == int(new_next_1c[b]) for b in range(args.batch)] + tok_ok = bool((new_cont >= 0).all() and (new_cont < m.cfg.vocab_size).all()) + print(f"[gate1] first-token match (B vs A) per lane: {first_match}") + print(f"[gate1] {args.cont}-token continuation match per lane: {lane_match}") + print(f"[gate1] chunked({args.chunk}) vs single-chunk next-token match: {chunk_vs_1c}") + print(f"[gate1] all tokens in-vocab: {tok_ok}", flush=True) + for b in range(args.batch): + if not lane_match[b]: + a, c = base_cont[b].tolist(), new_cont[b].tolist() + d = next(i for i in range(len(a)) if a[i] != c[i]) + print(f"[gate1] lane {b} diverges at cont idx {d}: A={a[max(0,d-2):d+3]} B={c[max(0,d-2):d+3]}") + # gate 5: decode aggregate tok/s after prefill (same lockstep path both sides) + dec_a = args.batch * (args.cont - 1) / dt_dec_a + dec_b = args.batch * (args.cont - 1) / dt_dec_b + print(f"[gate5] decode aggregate after A: {dec_a:.2f} tok/s; after B: {dec_b:.2f} tok/s; " + f"ratio B/A={dec_b/dec_a:.3f}", flush=True) + results["gate1"] = {"first_match": first_match, "cont_match": lane_match, + "chunk_vs_single": chunk_vs_1c, "tokens_in_vocab": tok_ok, + "base_next": [int(x) for x in base_next], + "new_next": [int(x) for x in new_next], + "base_cont": base_cont.tolist(), "new_cont": new_cont.tolist()} + results["gate5"] = {"dec_a_tok_s": dec_a, "dec_b_tok_s": dec_b, "ratio": dec_b / dec_a} + + ids_same = np.tile(ids[0], (args.batch, 1)) + m.reset() + same_next = m.prefill_batched(ids_same, chunk_size=args.chunk) + print(f"[gate1b] identical prompts -> identical next tokens: " + f"{bool((same_next == same_next[0]).all())} ({same_next.tolist()})", flush=True) + results["gate1b_same_prompt_next"] = [int(x) for x in same_next] + +# ---- Gate 4: TTFT A/B -------------------------------------------------------- +# One interleaved triple per rep (A tbt, B single-chunk, B chunk=N) so the slow +# A side (T lockstep steps) is measured once and both B configs share its +# thermal/contention window. +def ttft_ab(T, chunk, reps): + p = make_prompts(T) + for _ in range(2): # warmups + m.reset(); prefill_tbt(p); time.sleep(0.3) + m.reset(); m.prefill_batched(p, chunk_size=0); time.sleep(0.3) + m.reset(); m.prefill_batched(p, chunk_size=chunk); time.sleep(0.3) + a, b1, bc = [], [], [] + for r in range(reps): + m.reset(); t0 = time.perf_counter(); prefill_tbt(p) + a.append(time.perf_counter() - t0); time.sleep(0.3) + m.reset(); t0 = time.perf_counter(); m.prefill_batched(p, chunk_size=0) + b1.append(time.perf_counter() - t0); time.sleep(0.3) + m.reset(); t0 = time.perf_counter(); m.prefill_batched(p, chunk_size=chunk) + bc.append(time.perf_counter() - t0); time.sleep(0.3) + print(f"[ttft rep {r}] A={a[-1]*1e3:.1f}ms B(1chunk)={b1[-1]*1e3:.1f}ms " + f"B(chunk{chunk})={bc[-1]*1e3:.1f}ms", flush=True) + return a, b1, bc + +ra, rb1, rbc = ttft_ab(args.T, args.chunk, args.reps) +results["ttft"] = [] +ma = statistics.median(ra) +for chunk, rb in ((0, rb1), (args.chunk, rbc)): + mb = statistics.median(rb) + sp = ma / mb + print(f"[ttft] T={args.T:4d} chunk={chunk or 'seq_max'}: A(tbt)={ma*1e3:8.1f}ms " + f"B(batched)={mb*1e3:8.1f}ms speedup={sp:5.2f}x " + f"improvement={(1-mb/ma)*100:5.1f}%", flush=True) + results["ttft"].append(dict(T=args.T, chunk=chunk, A_ms=ma*1e3, B_ms=mb*1e3, + speedup=sp, A_all=[x*1e3 for x in ra], + B_all=[x*1e3 for x in rb])) + +if args.json_out: + with open(args.json_out, "w") as f: + json.dump(results, f, indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/lockstep_identity.py b/temp/gemma4_bprefill/lockstep_identity.py new file mode 100644 index 0000000..7a442b3 --- /dev/null +++ b/temp/gemma4_bprefill/lockstep_identity.py @@ -0,0 +1,44 @@ +"""Gate 3 (lockstep, batch=8): forward_batched old-path token identity. + +Run once with SK_DYLIB= and once with SK_DYLIB=; JSONs must match. +T kept small (32) — this only proves the old lockstep path is untouched. +""" +import os, sys, time, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r2"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--batch", type=int, default=8) +ap.add_argument("--T", type=int, default=32) +ap.add_argument("--cont", type=int, default=16) +ap.add_argument("--window", type=int, default=0) # 0 = model default +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +over = dict(batch=args.batch, seq_max=128, cache_max=512) +if args.window: + over["window"] = args.window +t0 = time.perf_counter() +m = sk.load("gemma4-12b-unified", **over) +print(f"[setup] loaded in {time.perf_counter()-t0:.1f}s (batch={args.batch})", flush=True) + +rng = np.random.default_rng(11) +ids = np.ascontiguousarray(rng.integers(10, 200000, size=(args.batch, args.T)).astype(np.int32)) + +m.reset() +out = None +for t in range(args.T): + out = m.forward_batched(np.ascontiguousarray(ids[:, t])) +toks = [np.array(out, dtype=np.int32).copy()] +for _ in range(args.cont): + toks.append(m.forward_batched(toks[-1]).astype(np.int32).copy()) +seqs = np.stack(toks, axis=1) +print(f"[lockstep] per-lane tokens: {seqs.tolist()}", flush=True) + +results = {"dylib": os.environ.get("SK_DYLIB", "?"), "lockstep": seqs.tolist()} +if args.json_out: + with open(args.json_out, "w") as f: + json.dump(results, f, indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/prompts_fluent.py b/temp/gemma4_bprefill/prompts_fluent.py new file mode 100644 index 0000000..88e6380 --- /dev/null +++ b/temp/gemma4_bprefill/prompts_fluent.py @@ -0,0 +1,10 @@ +TEXTS = [ + "The history of bread baking stretches back over ten thousand years, beginning with flat unleavened cakes cooked on hot stones beside open fires. Ancient Egyptian bakers discovered that dough left to rest would rise on its own, captured wild yeasts turning a dense paste into an airy loaf. Roman bakeries industrialized the craft with large masonry ovens and professional guilds, and by the Middle Ages nearly every European village supported a communal oven where families brought their shaped loaves each morning. The chemistry behind all of this remained mysterious until the nineteenth century, when scientists finally described how yeast ferments sugars into carbon dioxide, stretching gluten networks that bakers had been kneading by intuition for millennia. Modern artisans now combine that scientific understanding with", + "Metal compute shaders allow a programmer to dispatch thousands of threadgroups across the GPU, each cooperating through fast threadgroup memory while reading and writing device buffers. A well designed kernel keeps its arithmetic intensity high enough to hide memory latency, streaming tiles of data through shared storage so that every byte fetched from DRAM is reused many times before being discarded. On Apple Silicon the unified memory architecture removes explicit copies between processor and accelerator, but the bandwidth ceiling still dominates performance for most inference workloads. Profiling therefore begins with a simple roofline analysis: count the bytes each kernel must move, divide by the measured bandwidth of the chip, and compare that bound against the observed execution time to decide whether", + "In the deep ocean, hydrothermal vents support entire ecosystems that never see sunlight, powered instead by chemosynthetic bacteria that oxidize hydrogen sulfide gushing from the seafloor. Giant tube worms cluster around the vents in dense thickets, lacking mouths and digestive tracts entirely, nourished by symbiotic microbes housed within their tissues. Crabs, shrimp, and pale octopuses patrol the mineral chimneys, grazing on bacterial mats that coat every surface. When a vent finally goes extinct, the community collapses within months, yet larvae carried on deep currents somehow locate new vents tens of kilometers away and rebuild the entire assemblage. Biologists studying these habitats argue that similar chemical gardens beneath the ice of distant moons could plausibly", + "A well tuned sourdough starter doubles in volume within four to six hours of feeding, producing a pleasant aroma of ripe fruit and yogurt rather than the sharp smell of acetone that signals neglect. Maintaining that vigor requires a steady rhythm: discard most of the culture, refresh it with equal weights of flour and water, and hold it at a temperature where the yeasts and lactic acid bacteria stay in balance. Bakers who keep their starter in the refrigerator slow this cycle to a weekly feeding, trading some liveliness for convenience. The microbial community inside a mature starter is remarkably stable, resisting invasion by stray organisms because its acidity and competitive ecology leave no niche unfilled, which explains why", + "The transformer architecture replaced recurrence with attention, letting every token attend directly to every other token in the sequence instead of squeezing history through a fixed size hidden state. This change unlocked massive parallelism during training, since whole sequences could be processed at once on modern accelerators rather than step by step. The cost is quadratic scaling in sequence length, which has driven a decade of research into sparse patterns, sliding windows, linear approximations, and key value caching strategies. During inference the bottleneck shifts: generating each new token requires reading the entire stack of cached keys and values, so memory bandwidth rather than raw arithmetic throughput usually determines how quickly a large language model can", + "Glaciers carve valleys over millennia, grinding bedrock into fine flour that turns meltwater lakes a striking shade of turquoise blue. As the ice advances it plucks boulders from the valley walls and drags them along its base, scouring deep U shaped troughs that remain long after the climate warms. Moraines of tumbled rock mark each pause in the retreat, and stranded blocks of ice leave kettle ponds dotting the outwash plain. Geologists read these landforms like a written record, reconstructing the extent of ancient ice sheets from ridgelines and erratic boulders perched far from their parent outcrops. Today laser altimetry and satellite gravimetry extend that record forward, measuring yearly losses of ice mass that", + "Compilers translate high level source code into machine instructions through stages of parsing, optimization, and code generation, each built on decades of formal theory. The front end checks syntax and types, lowering the program into an intermediate representation that captures its meaning while discarding surface details. Optimization passes then rewrite this representation, folding constants, eliminating dead branches, hoisting loop invariant work, and vectorizing inner loops where the hardware allows. Register allocation maps an unbounded supply of virtual values onto a handful of physical registers, spilling the overflow to the stack as cheaply as possible. The final emission stage schedules instructions to keep the processor pipelines full, because even a perfectly optimized sequence of operations will", + "The annual monsoon arrives on the southwest coast in early June, bringing weeks of heavy rain that replenish rivers and aquifers after the long dry season. Farmers time their sowing to the first reliable downpours, and an early or late onset can shift harvests across an entire subcontinent. The system itself is a vast heat engine: land warms faster than ocean in spring, drawing moist air inland where it rises over the mountains and releases its water. Forecasters track sea surface temperatures thousands of kilometers away because shifts in the Pacific can strengthen or starve the circulation months in advance. Reservoir managers, city planners, and insurance companies all build their yearly calendars around the", +] diff --git a/temp/gemma4_bprefill/seq_paths_base.py b/temp/gemma4_bprefill/seq_paths_base.py new file mode 100644 index 0000000..ae798f2 --- /dev/null +++ b/temp/gemma4_bprefill/seq_paths_base.py @@ -0,0 +1,58 @@ +"""Pre-existing-path divergence demo (run on the BASE dylib, batch=1). + +Compares two EXISTING prefill paths for the same prompt: + path1 = one forward(T) call (seq>1 'original' attention path) + path2 = T x forward(seq=1) (decode fast-path attention) +then 32 greedy continuations each. Divergence here is inherited seq>1-vs-seq=1 +reduction-order numerics — present without any batched-prefill code. +""" +import os, sys, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r2"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--T", type=int, default=128) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--random", action="store_true") +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +m = sk.load("gemma4-12b-unified", seq_max=128, cache_max=512, window=512) +print(f"[setup] batch=1 window={m.cfg.window}", flush=True) + +if args.random: + rng = np.random.default_rng(7) + ids = rng.integers(10, 200000, size=args.T).astype(np.int32) +else: + txt = ("The history of bread baking stretches back over ten thousand years, " + "beginning with flat unleavened cakes cooked on hot stones. ") * 6 + ids = np.asarray(m.tokenizer.encode(txt)[:args.T], dtype=np.int32) + +def cont_from(first, n): + toks = [int(first)] + for _ in range(n - 1): + toks.append(int(m.forward(np.array([toks[-1]], dtype=np.int32))[0])) + return toks + +m.reset() +n1 = m.forward(ids) # one seq=T call +c1 = cont_from(n1[0], args.cont) + +m.reset() +n2 = None +for t in range(args.T): # T x seq=1 calls + n2 = m.forward(ids[t:t+1]) +c2 = cont_from(n2[0], args.cont) + +match = c1 == c2 +print(f"[seqpaths] first token equal: {int(n1[0]) == int(n2[0])}") +print(f"[seqpaths] {args.cont}-token continuation equal: {match}") +if not match: + d = next(i for i in range(len(c1)) if c1[i] != c2[i]) + print(f"[seqpaths] diverges at idx {d}: seqT={c1[max(0,d-2):d+3]} tbt={c2[max(0,d-2):d+3]}") +if args.json_out: + json.dump(dict(seqT_first=int(n1[0]), tbt_first=int(n2[0]), seqT_cont=c1, tbt_cont=c2), + open(args.json_out, "w"), indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/seq_paths_chat.py b/temp/gemma4_bprefill/seq_paths_chat.py new file mode 100644 index 0000000..a613f52 --- /dev/null +++ b/temp/gemma4_bprefill/seq_paths_chat.py @@ -0,0 +1,63 @@ +"""Inherited-divergence demo, CHAT-TEMPLATED prompts, BASE dylib (batch=1, zero +new code). Same spliced-to-T prompts as gate_chat.py: path1 = one forward(T) +(q_seq>1 attention), path2 = T x forward(seq=1) (decode fast path), 32 greedy +continuations each. path1's next token is also the single-stream seq>1 +reference for the batched gate's per-lane first token (gate_chat.json). +""" +import os, sys, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r4"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--T", type=int, default=128) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +m = sk.load("gemma4-12b-unified", seq_max=128, cache_max=512, window=512) +print(f"[setup] batch=1 window={m.cfg.window} dylib={os.environ.get('SK_DYLIB','?')}", + flush=True) + +from prompts_fluent import TEXTS + +def chat_ids(text): + return list(m.tokenizer.chat( + [{"role": "user", "content": "Continue this passage:\n\n" + text}], bos=True)) + +a, b = chat_ids(TEXTS[0]), chat_ids(TEXTS[1]) +S = 0 +while S < min(len(a), len(b)) and a[-1 - S] == b[-1 - S]: + S += 1 + +def cont32(n): + toks = [int(n[0])] + for _ in range(args.cont - 1): + toks.append(int(m.forward(np.array([toks[-1]], dtype=np.int32))[0])) + return toks + +results = [] +for i, t in enumerate(TEXTS): + full = chat_ids(t) + ids = np.asarray(full[: args.T - S] + full[-S:], dtype=np.int32) + m.reset() + n1 = m.forward(ids) + c1 = cont32(n1) + m.reset() + n2 = None + for k in range(len(ids)): + n2 = m.forward(ids[k:k+1]) + c2 = cont32(n2) + match = c1 == c2 + d = next((j for j in range(args.cont) if c1[j] != c2[j]), -1) + print(f"[spc] prompt {i} match={match} first_div_idx={d} next_seqT={c1[0]}", + flush=True) + results.append(dict(prompt=i, match=match, div_idx=d, c1=c1, c2=c2)) + +n_match = sum(r["match"] for r in results) +print(f"[spc] forward(T) vs T x forward(1), chat prompts: {n_match}/8 32-tok matches", + flush=True) +if args.json_out: + json.dump(results, open(args.json_out, "w"), indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/seq_paths_fluent.py b/temp/gemma4_bprefill/seq_paths_fluent.py new file mode 100644 index 0000000..4520bf5 --- /dev/null +++ b/temp/gemma4_bprefill/seq_paths_fluent.py @@ -0,0 +1,52 @@ +"""Inherited-divergence demo on the BASE dylib (batch=1, ZERO new code). + +For each fluent prompt: path1 = one forward(T) (q_seq>1 attention path), +path2 = T x forward(seq=1) (decode fast path), then 32 greedy continuations. +Any divergence here is pre-existing seq>1-vs-seq=1 reduction-order numerics — +the same numeric gap the batched-prefill A/B comparison crosses. +""" +import os, sys, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r4"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--T", type=int, default=120) +ap.add_argument("--cont", type=int, default=32) +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +m = sk.load("gemma4-12b-unified", seq_max=128, cache_max=512, window=512) +print(f"[setup] batch=1 window={m.cfg.window} dylib={os.environ.get('SK_DYLIB','?')}", + flush=True) + +from prompts_fluent import TEXTS + +def cont32(n): + toks = [int(n[0])] + for _ in range(args.cont - 1): + toks.append(int(m.forward(np.array([toks[-1]], dtype=np.int32))[0])) + return toks + +results = [] +for i, t in enumerate(TEXTS): + ids = np.asarray(m.tokenizer.encode(t)[: args.T], dtype=np.int32) + m.reset() + n1 = m.forward(ids) # one seq>1 chunk + c1 = cont32(n1) + m.reset() + n2 = None + for k in range(len(ids)): # token-by-token + n2 = m.forward(ids[k:k+1]) + c2 = cont32(n2) + match = c1 == c2 + d = next((j for j in range(args.cont) if c1[j] != c2[j]), -1) + print(f"[sp] prompt {i} match={match} first_div_idx={d}", flush=True) + results.append(dict(prompt=i, match=match, div_idx=d, c1=c1, c2=c2)) + +n_match = sum(r["match"] for r in results) +print(f"[sp] forward(T) vs T x forward(1): {n_match}/8 32-tok matches", flush=True) +if args.json_out: + json.dump(results, open(args.json_out, "w"), indent=1) +print("[done]", flush=True) diff --git a/temp/gemma4_bprefill/single_stream.py b/temp/gemma4_bprefill/single_stream.py new file mode 100644 index 0000000..a6210f6 --- /dev/null +++ b/temp/gemma4_bprefill/single_stream.py @@ -0,0 +1,47 @@ +"""Gate 0/3 (single-stream, batch=1): 12B coherence + old-path token identity. + +Run once with SK_DYLIB= and once with SK_DYLIB=; JSONs must match +token-for-token (the patched paths are additive: batched_prefill defaults 0). +""" +import os, sys, time, json, argparse + +sys.path.insert(0, os.path.expanduser(os.environ.get("SK_TREE", "~/sk-gemma-bprefill-r2"))) +import numpy as np +import SuperKittens as sk + +ap = argparse.ArgumentParser() +ap.add_argument("--json_out", default="") +args = ap.parse_args() + +t0 = time.perf_counter() +m = sk.load("gemma4-12b-unified", seq_max=128, cache_max=512) +print(f"[setup] loaded in {time.perf_counter()-t0:.1f}s (batch=1)", flush=True) + +results = {"dylib": os.environ.get("SK_DYLIB", "?")} + +# Gate 0: coherence (greedy, 32+ tokens). +ids = np.asarray(m.tokenizer.chat([{"role": "user", + "content": "Generate a poem about pizza dough"}], + bos=True), dtype=np.int32) +m.reset() +out = m.generate(ids, max_new_tokens=48, temperature=0.0) +text = m.tokenizer.decode(out, skip_special=True) +print(f"[coherence] {text!r}", flush=True) +results["coherence_ids"] = [int(x) for x in out] +results["coherence_text"] = text + +# Plain forward + 16 decode steps from a fixed token prompt. +pids = np.asarray(m.tokenizer.encode("The capital of France is"), dtype=np.int32) +m.reset() +nxt = m.forward(pids) +dec = [int(nxt[0])] +for _ in range(16): + nxt = m.forward(np.array([dec[-1]], dtype=np.int32)) + dec.append(int(nxt[0])) +print(f"[decode_ids] {dec}", flush=True) +results["decode_ids"] = dec + +if args.json_out: + with open(args.json_out, "w") as f: + json.dump(results, f, indent=1) +print("[done]", flush=True) From 6a108606fa0e5201f0252ced3cba52b6f2898a74 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 17:55:52 -0400 Subject: [PATCH 06/31] =?UTF-8?q?diffgemma:=20Phase=20A=20blueprint=20?= =?UTF-8?q?=E2=80=94=20arch=20contract=20(dual=20head=5Fdim=20ISWA=20MoE,?= =?UTF-8?q?=20bidirectional=20canvas),=20sampler=20spec=20(entropy-bound?= =?UTF-8?q?=20+=20SC),=20byte=20budget=20(experts=2014.09/15.64=20GiB),=20?= =?UTF-8?q?2-host=20fit=20plan,=20staged=20port?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- temp/diffgemma_feas/STATUS.md | 121 ++++++++++++++++++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 temp/diffgemma_feas/STATUS.md diff --git a/temp/diffgemma_feas/STATUS.md b/temp/diffgemma_feas/STATUS.md new file mode 100644 index 0000000..4d86e5d --- /dev/null +++ b/temp/diffgemma_feas/STATUS.md @@ -0,0 +1,121 @@ +# DiffusionGemma-26B-A4B-it → SK port blueprint (Phase A) + +Sources: unsloth/diffusiongemma-26B-A4B-it-GGUF Q4_K_M (downloaded, amelia +~/diffgemma-gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf, 16,806,810,336 B; +metadata dump at ~/diffgemma-gguf/dump.txt) + llama.cpp draft PR +ggml-org#24423 (full diff at /tmp/diffgemma_pr.diff on the laptop; key files +examples/diffusion/diffusion.cpp, src/models/diffusion-gemma.cpp). + +## Architecture (verified from GGUF header) + +| field | value | +|---|---| +| arch | `diffusion-gemma`, 30 layers, d_model 2816, vocab 262144 | +| ISWA | pattern [SWA×5, global]×6 (`sliding_window_pattern`), window 1024 | +| heads | 16 q-heads all layers; kv-heads 8 (SWA) / 2 (global) | +| head_dim | **256 (SWA) / 512 (global)** — dual dims; rope dims match (256/512) | +| rope | theta 1e6 (global) / 1e4 (SWA) | +| MoE | 128 routed experts (ff 704) + dense/shared ffn 2112, 8 used; expert tensors Q4_K/Q6_K/Q8_0 mix | +| head | tied embed (no separate output tensor beyond 0.56GiB embed), final softcap 30.0 | +| attention.causal | **false** | +| diffusion | canvas_length 256, mask_token_id 4 (``) | + +Byte budget (Q4_K_M, total 15.64 GiB): routed experts **14.09 GiB (90%)**; +attention 0.60; embed 0.56; dense ffn 0.37; norms/scales+head ~0.02. +Non-expert backbone ≈ **1.55 GiB**. + +## Reference runtime contract (from PR #24423) + +Graph (src/models/diffusion-gemma.cpp): +- Backbone "identical to gemma4" (shared weights; gemma4-common). ONE unified + forward over [prompt | canvas], split P = n_tokens − C. Region-aware: + 1. embeddings: prompt = embed·sqrt(n_embd); canvas = rmsnorm_noscale(embed·sqrt(n_embd)) + 2. per-layer scalars: `blk.N.enc_layer_output_scale` (prompt rows) vs + `blk.N.layer_output_scale` (canvas rows) — tiny F32 scalars + 3. additive mask: prompt queries causal over prompt only (SWA-clipped); + canvas queries bidirectional — global layers see ALL prompt+canvas, + SWA layers see last (n_swa−1) prompt positions + all canvas +- "A single no-cache bidirectional forward over [prompt|canvas] reproduces + the two-pass (causal encoder prefill + bidirectional decoder, zero + self-conditioning) result" → Stage-1 needs NO KV machinery. +- Cached mode (perf): prefill prompt once into prefix KV; per step decode + canvas-only with rectangular mask [P+C, C]. +- Self-conditioning (SC): previous step's RAW canvas logits [n_vocab, C] + uploaded as an input; SC subgraph feeds canvas embedding (gated off at + step 0 via scale 0; uses softmax at prev step's 1/t). + +Sampler (examples/diffusion/diffusion.cpp, EntropyBoundSampler): +- Canvas RANDOM-initialized (uniform over vocab, NOT mask tokens). +- Loop cur_step = S..1: t = t_min + (t_max−t_min)·(cur_step/S); per position: + argmax, entropy H of softmax(logits/t), one multinomial sample; stash raw + row for SC. +- Accept the lowest-entropy positions while the cumulative entropy of + strictly-earlier accepted positions ≤ entropy_bound (MI budget); accepted + positions take their SAMPLED token, the rest are RE-RANDOMIZED. +- Output = argmax canvas (not the working canvas). Adaptive stop: argmax + unchanged for stability_threshold steps AND mean entropy < + confidence_threshold. suppress_mask_token: logit[mask]=−inf. +- Defaults seen in PR flags: ~48 max steps, t 1.0→0.6. + +## Fit plan (16GB M4 minis, wired ceiling ~12.7 GiB) + +- Single-host RESIDENT: impossible (15.64 GiB). +- **2-host layer split (target)**: ~15 layers/host ≈ 7.0 GiB experts + share + of backbone → ~7.8 GiB resident/host. Comfortable. Needs full GGUF on both + hosts' disk (amelia has it; derek needs ~7 GiB freed — candidate + Qwen3-14B-Q4_K_M 8.5 GiB, USER AUTHORIZATION required). +- Single-host degraded (Stages 1–2 correctness): mmap + paging on amelia. + Expected expert touch per step: 8/128 per token × 256 tokens → essentially + all 128 experts per layer per step → full-file reads when paging; slow but + correct for small canvases/prompts. Use C=64 canvas + short prompts for + validation to bound the working set. +- Per-step estimate (2-host resident): expert-read bound ≈ 7 GiB/host / + ~110 GB/s ≈ 65 ms + attention/compute + hop → 150–500 ms/step realistic + band; 48 steps/256-tok block → **~10–35 tok/s 26B-class** (vs ~5 tok/s AR + if it fit). 3090 reference 326 ms/step at 8× bandwidth ⇒ temper to the low + end until measured; adaptive early-stop often cuts steps well below S. + +## SK port surface + +REUSES: gemma4 launcher skeleton (ISWA layout, rope pair, softcap, qk-norm), +deepseek MoE kernels (moe_group/mul_mv_id/down_scatter) + Q4_K MoE lab port, +gemm_mma (M = P+C rows ≈ prefill shapes), loader GGUF plumbing. + +NEW (layout: models/gemma/diffusion/ family package; kernels only there or +kernels//): +1. Loader rows for `diffusion-gemma` arch keys + enc/dec scale tensors. +2. Masked attention: additive-mask (or predicate) attention kernel for + head_dim 256 and 512 at M=P+C — the D=128 dense kernels do NOT apply; + simplest correct path = QK^T GEMM + mask add + softmax + V GEMM via + existing gemm_mma pieces (prefill-style, no streaming KV) for Stage 1. +3. Region-aware embedding + per-layer scalar plumbing. +4. Sampler host loop (Python first; mirrors the reference exactly, incl. + seed-reproducible RNG and SC buffer). +5. SC subgraph (softmax(prev logits/t_prev) → canvas embedding mix) — can be + STUBBED OFF for Stage 1 (reference gates it off step 0 / zero-SC unified + forward is the documented equivalence). +6. Stage 3: prefix-KV cached mode + 2-host layer-range split (pipeline.py + generalization for this family). + +## Stages + +- **Stage 1 — logits parity**: unified no-cache forward, zero SC, small + prompt + C=64..256 canvas on amelia (mmap OK). Validate per-position canvas + logits vs llama.cpp PR build (build llama-diffusion-cli on amelia from the + PR branch — llama.cpp runtime-compiles Metal, CLT-only OK; or CPU eval). + Gate: max rel err small + argmax canvas identical on ≥3 seeds/prompts. +- **Stage 2 — e2e coherent**: sampler loop + adaptive stop + chat template; + greedy/argmax output coherent; reproduce a reference generation + token-for-token at fixed seed (CPU RNG is ours = exact match possible). +- **Stage 3 — perf**: prefix-KV cached canvas decode; 2-host expert split; + measure ms/step + e2e tok/s; tune (expert batching at M=256 is + prefill-like — MMA grouped path). + +## Risks +1. Dual head_dim (256/512) attention — new territory; mitigated by + GEMM-composed attention for Stage 1 (correctness first). +2. SC subgraph semantics (exact mixing op) — extract precisely from the PR + model file before Stage 2; Stage 1 doesn't need it. +3. Paging thrash on single-host validation — bound with C=64 + short prompts; + colima resident on amelia shrinks headroom. +4. The PR is a moving draft — pin the commit SHA used for parity. From 61ee34ba22709e4cc9b0051cd5b15838e6e634fd Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 17:58:53 -0400 Subject: [PATCH 07/31] =?UTF-8?q?qwen=20prefill-attn=20profile:=20NO-GO=20?= =?UTF-8?q?=E2=80=94=20attention=20is=205.1%=20of=20TTFT@T=3D512=20(4B),?= =?UTF-8?q?=202.0%=20(14B);=20projections=20own=20prefill=20(3.31=20ms/tok?= =?UTF-8?q?=20linear=20term)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- temp/prof_prefill_attn/STATUS.md | 90 ++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 temp/prof_prefill_attn/STATUS.md diff --git a/temp/prof_prefill_attn/STATUS.md b/temp/prof_prefill_attn/STATUS.md new file mode 100644 index 0000000..966c75a --- /dev/null +++ b/temp/prof_prefill_attn/STATUS.md @@ -0,0 +1,90 @@ +# Prefill-attention share profile — qwen, single-stream TTFT (PROFILE-ONLY) + +**Question:** at realistic prompt lengths (T=512), what fraction of single-stream +prefill wall time is the attention stage (Q@K^T + softmax + @V over T×T) vs the +MMA projections — is a prefill-attention optimization a ≥10% e2e TTFT lever? + +**Answer: NO-GO. Attention is 5.1% of TTFT at T=512 on 4B (9.3% at T=1024, +2.0% on 14B at T=512). A hypothetical 2× faster prefill attention buys 2.6% +e2e at T=512 on 4B, 1.0% on 14B.** + +## Setup +- Host: derek (M4 base, 16 GB), CLT-only — clang++ dylib + `SK_METAL_SRC_FALLBACK` + runtime compile. Fresh `~/sk-prof-pfattn`, tree = local `main` @ 096c4d6. +- Model: `qwen3-4b-q4km` (`~/qwen-gguf/Qwen3-4B-Q4_K_M.gguf`), seq_max=1024, + cache_max=2048; 14B spot check seq_max=512, cache_max=1024. +- Method: single-chunk `sk_qwen_forward(T)` (one command buffer; stages not + interleaved). GPU-busy via `SK_QWEN_GPUPROF` (GPUEnd−GPUStart). 2 warmups + + 5 reps median, 0.3 s gaps. Rep spread ≤0.2% for T≥256 (≤1.7% overall). +- Attribution: stage-skip gates already in `models/qwen/qwen_model.h` + (`SK_PROF_SKIP_ATTN` drops only the `mha_causal_prefill` dispatch at seq>1; + `SK_PROF_SKIP_XPOSE` drops the 4 seq<->head transposes; decode untouched). + Separate process per mode (gates latch as static consts). +- NOTE: first chain attempt was poisoned by a concurrent gemma-12B gate run on + derek (T=32 read 20.5 s gpu-busy under swap thrash). Killed my chain, waited + for the box to clear, reran clean. All numbers below are from the clean run. + +## TTFT(T) sweep — qwen3-4B-Q4_K_M (gpu-busy median, ms) + +| T | base | skip_attn | attn (Δ) | attn share | skip_xpose Δ | +|------|---------|-----------|----------|------------|--------------| +| 32 | 145.2 | 145.6 | −0.3 | (noise) | +0.96 (0.7%) | +| 64 | 251.6 | 248.5 | 3.1 | 1.2% | +0.43 (0.2%) | +| 128 | 474.9 | 470.6 | 4.3 | 0.9% | +1.62 (0.3%) | +| 256 | 904.5 | 878.9 | 25.6 | 2.8% | +1.25 (0.1%) | +| 512 | 1849.8 | 1756.1 | 93.6 | **5.1%** | +4.29 (0.2%) | +| 1024 | 3870.3 | 3510.0 | 360.4 | **9.3%** | +11.9 (0.3%) | + +Wall ≈ gpu_busy + 2-4 ms at every T — prefill is fully GPU-bound; host overhead +is irrelevant. Transposes are free (skip_xpose Δ < 0.5% everywhere). + +## Fit: TTFT(T) = c + a·T + b·T² (relative-error-weighted LSQ on base gpu_med) + +c = 38.9 ms, a = 3.311 ms/tok, b = 0.4207 µs/tok². Fit error ≤1.1% at all six T. + +| T | quad share (fit) | ablation share | +|------|------------------|----------------| +| 128 | 1.5% | 0.9% | +| 256 | 3.0% | 2.8% | +| 512 | 6.0% | 5.1% | +| 1024 | 11.4% | 9.3% | + +Fit and ablation agree; the fit's quadratic term runs slightly high because it +also absorbs the KV-write growth and MMA tail effects. TTFT is dominated by the +linear term: 3.31 ms/token of projection/MLP GEMM work (≈277 tok/s prefill; +the projections are the compute-bound cost, exactly as the gemm_mma design +intends — weights are streamed once per chunk, so there is no bandwidth-bound +weight re-read to hide). + +## 14B spot check (T=128/512, gpu-busy median, ms) + +| T | base | skip_attn | attn (Δ) | attn share | +|-----|--------|-----------|----------|------------| +| 128 | 1636.9 | 1625.6 | 11.3 | 0.7% | +| 512 | 6549.8 | 6417.2 | 132.6 | **2.0%** | + +Attention share SHRINKS with model size: attention grew 1.42× from 4B→14B +(93.6→132.6 ms at T=512, ≈ the heads·layers ratio 40·40/32·36 = 1.39×) while +the projection term grew ~3.5× (≈ the param ratio). Bigger models are even +less attention-bound at prefill. + +## Ceiling math +- Attention-FREE TTFT(512) on 4B = 1756 ms vs 1850 ms → only **5.1%** better. +- 2× faster prefill attention: 0.5 × share → **2.6% e2e at T=512**, 4.7% at T=1024. +- For a 2× attention win to clear 10% e2e, attention share must be ≥20%. + From the fit, share(T)=0.20 at **T ≈ 2000 tokens** — beyond typical prompts + and the default seq_max envelope. Even there a 2× win is exactly at the 10% bar. + +## GO/NO-GO: **NO-GO** +Attention is 5% of single-stream prefill at T=512 and ~9% at T=1024; the +projections' linear term (3.31 ms/tok) owns TTFT. The BR=8 `mha_causal_prefill` +already amortizes K/V re-streaming well enough that prefill-attention work is a +<3% e2e lever at realistic prompt lengths. Prefill optimization effort should +target the MMA projection path (the a-term), not attention. + +## Artifacts +- `prof_prefill.py` (sweep driver, fd2-redirect gpuprof capture), `analyze.py` + (fit + ablation cross-check), `build_dylib.sh`, `run.sh` (derek env wrapper). +- Raw JSONs: `base_4b.json`, `skipattn_4b.json`, `skipxpose_4b.json`, + `base_14b.json`, `skipattn_14b.json` (5 reps each, wall + gpu). +- derek: `~/sk-prof-pfattn/` (tree, dylib, logs). From 8837178353e702c5ae961fae5a011f8674a2f1ad Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 18:26:54 -0400 Subject: [PATCH 08/31] diffgemma stage1: family package (GGUF-native loader, ggml-mirror CPU oracle, Metal unified forward), op tests green --- SuperKittens/inference/registry.py | 13 + .../models/gemma/diffusion/__init__.py | 10 + .../models/gemma/diffusion/adapter.py | 28 ++ SuperKittens/models/gemma/diffusion/config.py | 74 ++++ .../models/gemma/diffusion/dg_kernels.metal | 64 +++ .../models/gemma/diffusion/forward_metal.py | 400 ++++++++++++++++++ .../models/gemma/diffusion/gguf_io.py | 231 ++++++++++ .../models/gemma/diffusion/graph_ref.py | 235 ++++++++++ SuperKittens/models/gemma/diffusion/runner.py | 84 ++++ temp/diffgemma_s1/test_ops.py | 133 ++++++ 10 files changed, 1272 insertions(+) create mode 100644 SuperKittens/models/gemma/diffusion/__init__.py create mode 100644 SuperKittens/models/gemma/diffusion/adapter.py create mode 100644 SuperKittens/models/gemma/diffusion/config.py create mode 100644 SuperKittens/models/gemma/diffusion/dg_kernels.metal create mode 100644 SuperKittens/models/gemma/diffusion/forward_metal.py create mode 100644 SuperKittens/models/gemma/diffusion/gguf_io.py create mode 100644 SuperKittens/models/gemma/diffusion/graph_ref.py create mode 100644 SuperKittens/models/gemma/diffusion/runner.py create mode 100644 temp/diffgemma_s1/test_ops.py diff --git a/SuperKittens/inference/registry.py b/SuperKittens/inference/registry.py index c588547..28481db 100644 --- a/SuperKittens/inference/registry.py +++ b/SuperKittens/inference/registry.py @@ -402,6 +402,19 @@ class ModelSpec: head_dim=128, n_int=25600, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, tie_word_embeddings=0), ), + # DiffusionGemma 26B-A4B: block text-diffusion MoE on a gemma4 backbone + # (llama.cpp PR #24423 is the runtime reference). Stage-1 adapter exposes + # the unified zero-SC forward only; the entropy-bound sampler is Stage 2. + # Dims live in the GGUF metadata (config_from_gguf), not here. + "diffgemma-26b": ModelSpec( + family="diffgemma", + adapter="SuperKittens.models.gemma.diffusion.adapter:DiffusionGemma", + hf_repo="unsloth/diffusiongemma-26B-A4B-it-GGUF", + weight_dir="diffgemma-26b", + gguf_name="diffusiongemma-26B-A4B-it-Q4_K_M.gguf", + default_quant="q4_k_m", + tokenizer_family="gemma4", + ), } diff --git a/SuperKittens/models/gemma/diffusion/__init__.py b/SuperKittens/models/gemma/diffusion/__init__.py new file mode 100644 index 0000000..60a8736 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/__init__.py @@ -0,0 +1,10 @@ +"""DiffusionGemma (block text-diffusion MoE on a Gemma-4 backbone) family. + +Stage 1 (logits parity) ships: GGUF-native loader (`gguf_io`, `config`), the +ggml-mirror CPU oracle (`graph_ref`), and the Metal unified forward +(`forward_metal`). Sampler / cached decode are Stage 2+. +""" +from .config import DiffusionGemmaConfig, config_from_gguf +from .gguf_io import GGUFFile + +__all__ = ["DiffusionGemmaConfig", "config_from_gguf", "GGUFFile"] diff --git a/SuperKittens/models/gemma/diffusion/adapter.py b/SuperKittens/models/gemma/diffusion/adapter.py new file mode 100644 index 0000000..cb2eac1 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/adapter.py @@ -0,0 +1,28 @@ +# pyright: reportMissingImports=false +"""adapter.py — registry seam for the DiffusionGemma family. + +Stage-1 scope: loading + the unified zero-SC forward (logits parity). The +denoising sampler loop / generation entrypoint land in Stage 2, so this +adapter intentionally exposes `forward(ids, P)` and not `generate`. +""" +from __future__ import annotations + +from pathlib import Path + +from .config import config_from_gguf +from .gguf_io import GGUFFile + + +class DiffusionGemma: + @classmethod + def from_spec(cls, spec, **overrides): + from .forward_metal import DiffusionGemmaMetal + + sk_root = Path(__file__).resolve().parents[3] + snap = Path(overrides.pop("snapshot", None) + or (sk_root / "model_weights" / spec.weight_dir)) + gguf = overrides.pop("gguf", None) or (snap / spec.gguf_name) + if not Path(gguf).exists(): + raise FileNotFoundError(f"DiffusionGemma GGUF not found: {gguf}") + cfg = config_from_gguf(GGUFFile(str(gguf)).meta) + return DiffusionGemmaMetal(str(gguf), cfg) diff --git a/SuperKittens/models/gemma/diffusion/config.py b/SuperKittens/models/gemma/diffusion/config.py new file mode 100644 index 0000000..2e658ad --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/config.py @@ -0,0 +1,74 @@ +"""config.py — DiffusionGemma family config, mapped from `diffusion-gemma` GGUF +metadata (per-layer kv-head array, sliding_window_pattern, dual head/rope dims, +expert counts, canvas_length, softcap, mask token).""" +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class DiffusionGemmaConfig: + n_layers: int = 30 + d_model: int = 2816 + n_heads: int = 16 + n_kv_heads: tuple[int, ...] = () # per layer (8 SWA / 2 global) + is_swa: tuple[bool, ...] = () # sliding_window_pattern (True = SWA) + head_dim_swa: int = 256 + head_dim_global: int = 512 + rope_dims_swa: int = 256 + rope_dims_global: int = 512 + rope_base: float = 1e6 # global layers + rope_base_swa: float = 1e4 + window: int = 1024 # n_swa + n_ff: int = 2112 # dense (shared-expert) MLP + n_ff_exp: int = 704 + n_expert: int = 128 + n_expert_used: int = 8 + vocab_size: int = 262144 + eps: float = 1e-6 + final_logit_softcap: float = 30.0 + attn_scale: float = 1.0 # gemma4: no pre-attn scaling (qk-norm) + canvas_length: int = 256 + mask_token_id: int = 4 + bos_token_id: int = 2 + + def head_dim(self, il: int) -> int: + return self.head_dim_swa if self.is_swa[il] else self.head_dim_global + + def rope_params(self, il: int) -> tuple[float, bool]: + """(freq_base, uses_freq_factors) for layer il.""" + if self.is_swa[il]: + return self.rope_base_swa, False + return self.rope_base, True + + +def config_from_gguf(meta: dict) -> DiffusionGemmaConfig: + p = "diffusion-gemma." + c = DiffusionGemmaConfig( + n_layers=int(meta[p + "block_count"]), + d_model=int(meta[p + "embedding_length"]), + n_heads=int(meta[p + "attention.head_count"]), + n_kv_heads=tuple(int(v) for v in meta[p + "attention.head_count_kv"]), + is_swa=tuple(bool(v) for v in meta[p + "attention.sliding_window_pattern"]), + head_dim_swa=int(meta[p + "attention.key_length_swa"]), + head_dim_global=int(meta[p + "attention.key_length"]), + rope_dims_swa=int(meta[p + "rope.dimension_count_swa"]), + rope_dims_global=int(meta[p + "rope.dimension_count"]), + rope_base=float(meta[p + "rope.freq_base"]), + rope_base_swa=float(meta[p + "rope.freq_base_swa"]), + window=int(meta[p + "attention.sliding_window"]), + n_ff=int(meta[p + "feed_forward_length"]), + n_ff_exp=int(meta[p + "expert_feed_forward_length"]), + n_expert=int(meta[p + "expert_count"]), + n_expert_used=int(meta[p + "expert_used_count"]), + vocab_size=len(meta["tokenizer.ggml.tokens"]), + eps=float(meta[p + "attention.layer_norm_rms_epsilon"]), + final_logit_softcap=float(meta[p + "final_logit_softcapping"]), + canvas_length=int(meta["diffusion.canvas_length"]), + mask_token_id=int(meta["tokenizer.ggml.mask_token_id"]), + bos_token_id=int(meta["tokenizer.ggml.bos_token_id"]), + ) + assert meta[p + "attention.causal"] is False + assert int(meta[p + "attention.value_length"]) == c.head_dim_global + assert int(meta[p + "attention.value_length_swa"]) == c.head_dim_swa + return c diff --git a/SuperKittens/models/gemma/diffusion/dg_kernels.metal b/SuperKittens/models/gemma/diffusion/dg_kernels.metal new file mode 100644 index 0000000..ea7f537 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/dg_kernels.metal @@ -0,0 +1,64 @@ +// dg_kernels.metal — DiffusionGemma family kernels. +// +// dg_softmax_mask: masked row softmax for the GEMM-composed unified attention +// (QK^T tiles come from kernels/gemm/gemm_mma.metal at M = P+C). The D=128 +// production attention kernels don't apply at head_dim 256/512, so Stage 1 +// runs QK^T -> this kernel -> @V as three plain dispatches. +// +// S : fp16 [R, ncols], R = n_heads * n_tok (head-major), softmaxed in place +// mask : f32 [n_tok, ncols] additive (0 / -inf); row r uses mask row r % n_tok +// (one mask per layer, shared across heads; pad cols carry -inf) +// scale: kq scale applied before the mask (1.0 for this family — qk-norm) + +#include +using namespace metal; + +kernel void dg_softmax_mask( + device half *S [[buffer(0)]], + device const float *mask [[buffer(1)]], + constant uint &ncols [[buffer(2)]], + constant uint &ntok [[buffer(3)]], + constant float &scale [[buffer(4)]], + uint3 tgpig [[threadgroup_position_in_grid]], + uint3 tid3 [[thread_position_in_threadgroup]], + uint3 tptg3 [[threads_per_threadgroup]]) +{ + const uint tid = tid3.x; + const uint tptg = tptg3.x; + const uint r = tgpig.x; + device half *row = S + (size_t)r * ncols; + device const float *mrow = mask + (size_t)(r % ntok) * ncols; + + threadgroup float red[256]; + + float mx = -INFINITY; + for (uint c = tid; c < ncols; c += tptg) { + mx = max(mx, (float)row[c] * scale + mrow[c]); + } + red[tid] = mx; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = tptg / 2; s > 0; s >>= 1) { + if (tid < s) red[tid] = max(red[tid], red[tid + s]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float rmax = red[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float sum = 0.0f; + for (uint c = tid; c < ncols; c += tptg) { + sum += exp((float)row[c] * scale + mrow[c] - rmax); + } + red[tid] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = tptg / 2; s > 0; s >>= 1) { + if (tid < s) red[tid] += red[tid + s]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float inv = 1.0f / red[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + // recompute exp from the original score so each prob is rounded to fp16 once + for (uint c = tid; c < ncols; c += tptg) { + row[c] = (half)(exp((float)row[c] * scale + mrow[c] - rmax) * inv); + } +} diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py new file mode 100644 index 0000000..f4d3ff9 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -0,0 +1,400 @@ +# pyright: reportAttributeAccessIssue=false, reportMissingImports=false +"""forward_metal.py — DiffusionGemma unified [prompt|canvas] forward on Metal. + +Stage-1 correctness driver. Weight matmuls run on SK quant GEMM kernels +(kernels/gemm/gemm_mma.metal, native Q4_K/Q6_K/Q8_0 off the mmap'd GGUF — +no-copy MTLBuffers, the OS pager is the streaming layer). Attention is +GEMM-composed: QK^T (gemm_mma_f16) -> dg_softmax_mask (additive region mask) +-> @V (gemm_mma_f16). Host glue (norms, rope, router, geglu, residuals, +region scalars) is f32 numpy via graph_ref — shared with the CPU oracle so a +layer-dump diff bisects GPU vs host exactly. + +Known Stage-1 perf debt (deliberate): per-op CPU round trips, per-expert GEMM +loop, fp16 activation casts at each hop. Stage 3 moves the glue on-device. +""" +from __future__ import annotations + +import math +import mmap +import os +from pathlib import Path + +import numpy as np +import objc # noqa: F401 +import Metal # type: ignore[import-not-found] + +from .config import DiffusionGemmaConfig +from .gguf_io import GGUFFile, TensorInfo +from .graph_ref import (F32, Weights, build_mask, embed_tokens, gelu_tanh, + moe_route, rms_norm, rope_neox) + +_SK_ROOT = Path(__file__).resolve().parents[3] +_KERNEL_SOURCES = [ + _SK_ROOT / "kernels" / "gemm" / "gemm_mma.metal", + Path(__file__).parent / "dg_kernels.metal", +] +_MMA_BY_TYPE = {"F16": "gemm_mma_f16", "BF16": "gemm_mma_bf16", + "Q8_0": "gemm_mma_q8_0", "Q4_K": "gemm_mma_q4k", + "Q6_K": "gemm_mma_q6k"} + + +def _pad32(n: int) -> int: + return (n + 31) // 32 * 32 + + +class MetalCtx: + """Device + queue + runtime-compiled PSOs (CLT-only hosts: no metallib).""" + + def __init__(self): + dev = Metal.MTLCreateSystemDefaultDevice() + if dev is None: + raise RuntimeError("no Metal device") + self.device = dev + self.queue = dev.newCommandQueue() + src = "\n".join(p.read_text() for p in _KERNEL_SOURCES) + opts = Metal.MTLCompileOptions.alloc().init() + # bfloat (gemm_mma_bf16) needs Metal >= 3.1 on runtime compile + opts.setLanguageVersion_(getattr(Metal, "MTLLanguageVersion3_1", (3 << 16) + 1)) + lib, err = dev.newLibraryWithSource_options_error_(src, opts, None) + if lib is None: + raise RuntimeError(f"metal compile failed: {err}") + self.lib = lib + self._pso = {} + + def pso(self, name: str): + p = self._pso.get(name) + if p is None: + fn = self.lib.newFunctionWithName_(name) + if fn is None: + raise RuntimeError(f"kernel not found: {name}") + p, err = self.device.newComputePipelineStateWithFunction_error_(fn, None) + if p is None: + raise RuntimeError(f"PSO failed for {name}: {err}") + self._pso[name] = p + return p + + def buf_from(self, arr: np.ndarray): + arr = np.ascontiguousarray(arr) + b = self.device.newBufferWithBytes_length_options_( + arr, arr.nbytes, Metal.MTLResourceStorageModeShared) + if b is None: + raise RuntimeError(f"buffer alloc failed ({arr.nbytes} B)") + return b + + def buf_empty(self, nbytes: int): + b = self.device.newBufferWithLength_options_( + nbytes, Metal.MTLResourceStorageModeShared) + if b is None: + raise RuntimeError(f"buffer alloc failed ({nbytes} B)") + return b + + def read(self, buf, dtype, shape, offset: int = 0) -> np.ndarray: + n = int(np.prod(shape)) * np.dtype(dtype).itemsize + mv = buf.contents().as_buffer(offset + n)[offset:offset + n] + return np.frombuffer(mv, dtype=dtype).reshape(shape).copy() + + +class WeightBufs: + """Per-tensor MTLBuffers over the GGUF. No-copy mmap windows when the + bridge cooperates (verified per process at init), else lazy copies.""" + + def __init__(self, ctx: MetalCtx, gg: GGUFFile): + self.ctx = ctx + self.gg = gg + self._cache: dict[str, tuple] = {} + self._f = open(gg.path, "rb") + self.nocopy_ok = self._probe_nocopy() + + def _map(self, ti: TensorInfo): + page = mmap.ALLOCATIONGRANULARITY + base = ti.offset // page * page + delta = ti.offset - base + length = (delta + ti.nbytes + page - 1) // page * page + # MAP_PRIVATE + writable prot: the bridge requires a writable buffer + # object for the void* arg; pages stay clean (we never write). + mm = mmap.mmap(self._f.fileno(), length, flags=mmap.MAP_PRIVATE, + prot=mmap.PROT_READ | mmap.PROT_WRITE, offset=base) + buf = self.ctx.device.newBufferWithBytesNoCopy_length_options_deallocator_( + mm, length, Metal.MTLResourceStorageModeShared, None) + if buf is None: + raise RuntimeError("newBufferWithBytesNoCopy returned None") + return buf, delta, mm + + def _probe_nocopy(self) -> bool: + try: + ti = min(self.gg.tensors.values(), key=lambda t: t.nbytes) + buf, delta, mm = self._map(ti) + got = bytes(buf.contents().as_buffer(delta + 16)[delta:delta + 16]) + want = bytes(self.gg.mm[ti.offset:ti.offset + 16]) + return got == want + except Exception as e: # noqa: BLE001 + print(f"[weights] no-copy probe failed ({e}); falling back to copies") + return False + + def get(self, name: str) -> tuple: + """-> (MTLBuffer, byte_offset, TensorInfo)""" + hit = self._cache.get(name) + if hit is not None: + return hit + ti = self.gg.tensors[name] + if self.nocopy_ok: + buf, delta, mm = self._map(ti) + ent = (buf, delta, ti, mm) + else: + arr = np.frombuffer(self.gg.mm, dtype=np.uint8, count=ti.nbytes, + offset=ti.offset) + ent = (self.ctx.buf_from(arr), 0, ti, None) + self._cache[name] = ent + return ent + + +class GemmBatch: + """Record gemm_mma / dg_softmax_mask dispatches into one command buffer; + memory barriers split dependent stages. run() commits + waits.""" + + def __init__(self, ctx: MetalCtx): + self.ctx = ctx + self.cmd = ctx.queue.commandBuffer() + self.enc = self.cmd.computeCommandEncoder() + self._keep: list = [] # encoder retains resources, but keep pyobjc refs too + self._u32 = lambda v: np.uint32(v).tobytes() + + def gemm(self, kernel: str, a_buf, a_off: int, w_buf, w_off: int, + c_buf, c_off: int, M: int, N: int, K: int, ldc: int | None = None): + ldc = N if ldc is None else ldc + enc = self.enc + self._keep += [a_buf, w_buf, c_buf] + enc.setComputePipelineState_(self.ctx.pso(kernel)) + enc.setBuffer_offset_atIndex_(a_buf, a_off, 0) + enc.setBuffer_offset_atIndex_(w_buf, w_off, 1) + enc.setBuffer_offset_atIndex_(c_buf, c_off, 2) + enc.setBytes_length_atIndex_(self._u32(M), 4, 3) + enc.setBytes_length_atIndex_(self._u32(N), 4, 4) + enc.setBytes_length_atIndex_(self._u32(K), 4, 5) + enc.setBytes_length_atIndex_(self._u32(ldc), 4, 6) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + Metal.MTLSizeMake((N + 31) // 32, (M + 31) // 32, 1), + Metal.MTLSizeMake(64, 1, 1)) + + def softmax_mask(self, s_buf, mask_buf, rows: int, ncols: int, ntok: int, + scale: float = 1.0): + enc = self.enc + self._keep += [s_buf, mask_buf] + enc.setComputePipelineState_(self.ctx.pso("dg_softmax_mask")) + enc.setBuffer_offset_atIndex_(s_buf, 0, 0) + enc.setBuffer_offset_atIndex_(mask_buf, 0, 1) + enc.setBytes_length_atIndex_(self._u32(ncols), 4, 2) + enc.setBytes_length_atIndex_(self._u32(ntok), 4, 3) + enc.setBytes_length_atIndex_(np.float32(scale).tobytes(), 4, 4) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + Metal.MTLSizeMake(rows, 1, 1), Metal.MTLSizeMake(256, 1, 1)) + + def barrier(self): + self.enc.memoryBarrierWithScope_(Metal.MTLBarrierScopeBuffers) + + def run(self): + self.enc.endEncoding() + self.cmd.commit() + self.cmd.waitUntilCompleted() + if self.cmd.error() is not None: + raise RuntimeError(f"command buffer failed: {self.cmd.error()}") + + +class DiffusionGemmaMetal: + """Stage-1 unified zero-SC forward. forward(ids, P) -> canvas logits f32.""" + + def __init__(self, gguf_path: str, cfg: DiffusionGemmaConfig): + self.gg = GGUFFile(gguf_path) + self.cfg = cfg + self.ctx = MetalCtx() + self.wb = WeightBufs(self.ctx, self.gg) + self.w = Weights(self.gg) # F32 sidecars (norms, scales, router) + self.dump = None # optional (name, il, arr) tap + + # -- helpers --------------------------------------------------------------- + + def _wgemm(self, batch: GemmBatch, wname: str, a_buf, a_off: int, + c_buf, c_off: int, M: int, N: int, K: int, + row0: int = 0, ldc: int | None = None): + """C[M,N] = A[M,K] @ W[row0:row0+N, :K]^T for GGUF weight `wname`.""" + buf, delta, ti, _ = self.wb.get(wname) + kern = _MMA_BY_TYPE[ti.type_name] + assert ti.shape[0] == K, (wname, ti.shape, K) + w_off = delta + row0 * ti.row_bytes + batch.gemm(kern, a_buf, a_off, buf, w_off, c_buf, c_off, M, N, K, ldc) + + def _gemm_f32(self, wname: str, a: np.ndarray, N: int, row0: int = 0) -> np.ndarray: + """One-shot weight GEMM with f32 host I/O (fp16 on the wire).""" + M, K = a.shape + amax = float(np.abs(a).max(initial=0.0)) + if amax > 3.0e4: + print(f"[warn] fp16 activation near overflow ({amax:.1f}) into {wname}") + a_buf = self.ctx.buf_from(a.astype(np.float16)) + c_buf = self.ctx.buf_empty(M * N * 2) + b = GemmBatch(self.ctx) + self._wgemm(b, wname, a_buf, 0, c_buf, 0, M, N, K, row0=row0) + b.run() + return self.ctx.read(c_buf, np.float16, (M, N)).astype(F32) + + # -- attention ------------------------------------------------------------- + + def _attention(self, il: int, q: np.ndarray, k: np.ndarray, v: np.ndarray, + mask: np.ndarray) -> np.ndarray: + """q [N,H,hd], k/v [N,Kv,hd] (post norm+rope, f32) -> [N, H*hd] f32.""" + cfg = self.cfg + N, H, hd = q.shape + Kv = k.shape[1] + gqa = H // Kv + Np = _pad32(N) + + qp = np.ascontiguousarray(q.transpose(1, 0, 2)).astype(np.float16) # [H,N,hd] + kp = np.zeros((Kv, Np, hd), np.float16) + kp[:, :N] = k.transpose(1, 0, 2) + vtp = np.zeros((Kv, hd, Np), np.float16) + vtp[:, :, :N] = v.transpose(1, 2, 0) + + q_buf = self.ctx.buf_from(qp) + k_buf = self.ctx.buf_from(kp) + vt_buf = self.ctx.buf_from(vtp) + m_buf = self.ctx.buf_from(np.ascontiguousarray(mask[:, :Np])) + s_buf = self.ctx.buf_empty(H * N * Np * 2) + o_buf = self.ctx.buf_empty(H * N * hd * 2) + + b = GemmBatch(self.ctx) + for h in range(H): + g = h // gqa + b.gemm("gemm_mma_f16", q_buf, h * N * hd * 2, k_buf, g * Np * hd * 2, + s_buf, h * N * Np * 2, M=N, N=Np, K=hd, ldc=Np) + b.barrier() + b.softmax_mask(s_buf, m_buf, rows=H * N, ncols=Np, ntok=N, + scale=cfg.attn_scale) + b.barrier() + for h in range(H): + g = h // gqa + b.gemm("gemm_mma_f16", s_buf, h * N * Np * 2, vt_buf, g * hd * Np * 2, + o_buf, h * N * hd * 2, M=N, N=hd, K=Np, ldc=hd) + b.run() + + o = self.ctx.read(o_buf, np.float16, (H, N, hd)).astype(F32) + return np.ascontiguousarray(o.transpose(1, 0, 2)).reshape(N, H * hd) + + # -- MoE ------------------------------------------------------------------- + + def _moe(self, il: int, attn_out: np.ndarray, e_in: np.ndarray) -> np.ndarray: + cfg = self.cfg + N = attn_out.shape[0] + sel, wts = moe_route(self.w, cfg, attn_out, il) + if self.dump: + self.dump("moe_sel", il, sel.astype(np.int32)) + self.dump("moe_wts", il, wts) + down_s = self.w.f32(f"blk.{il}.ffn_down_exps.scale") + gu_name = f"blk.{il}.ffn_gate_up_exps.weight" + dn_name = f"blk.{il}.ffn_down_exps.weight" + + experts = [] + for e in np.unique(sel): + tok, slot = np.nonzero(sel == e) + experts.append((int(e), tok, slot)) + + # stage A: gate_up for every hit expert in one command buffer + b = GemmBatch(self.ctx) + stage = [] + for e, tok, slot in experts: + m = len(tok) + a_buf = self.ctx.buf_from(e_in[tok].astype(np.float16)) + c_buf = self.ctx.buf_empty(m * 2 * cfg.n_ff_exp * 2) + self._wgemm(b, gu_name, a_buf, 0, c_buf, 0, M=m, N=2 * cfg.n_ff_exp, + K=cfg.d_model, row0=e * 2 * cfg.n_ff_exp) + stage.append((e, tok, slot, c_buf, m)) + b.run() + + # host geglu, then stage B: down for every hit expert + b = GemmBatch(self.ctx) + stage2 = [] + for e, tok, slot, c_buf, m in stage: + gu = self.ctx.read(c_buf, np.float16, (m, 2 * cfg.n_ff_exp)).astype(F32) + act = gelu_tanh(gu[:, :cfg.n_ff_exp]) * gu[:, cfg.n_ff_exp:] + a_buf = self.ctx.buf_from(act.astype(np.float16)) + d_buf = self.ctx.buf_empty(m * cfg.d_model * 2) + self._wgemm(b, dn_name, a_buf, 0, d_buf, 0, M=m, N=cfg.d_model, + K=cfg.n_ff_exp, row0=e * cfg.d_model) + stage2.append((e, tok, slot, d_buf, m)) + b.run() + + moe = np.zeros((N, cfg.d_model), dtype=F32) + for e, tok, slot, d_buf, m in stage2: + d_ = self.ctx.read(d_buf, np.float16, (m, cfg.d_model)).astype(F32) + moe[tok] += d_ * down_s[e] * wts[tok, slot][:, None] + return moe + + # -- forward --------------------------------------------------------------- + + def forward(self, ids: np.ndarray, P: int) -> np.ndarray: + cfg, w = self.cfg, self.w + ids = np.asarray(ids) + N = len(ids) + C = N - P + pos = np.arange(N, dtype=np.int64) + dmp = self.dump or (lambda name, il, arr: None) + + x = embed_tokens(w, cfg, ids, P) + dmp("inp_region", -1, x) + rope_ff = w.f32("rope_freqs.weight") + + for il in range(cfg.n_layers): + swa = cfg.is_swa[il] + hd = cfg.head_dim(il) + n_kv = cfg.n_kv_heads[il] + base, use_ff = cfg.rope_params(il) + ff = rope_ff if use_ff else None + + h = rms_norm(x, cfg.eps, w.f32(f"blk.{il}.attn_norm.weight")) + + q = self._gemm_f32(f"blk.{il}.attn_q.weight", h, cfg.n_heads * hd) + k_raw = self._gemm_f32(f"blk.{il}.attn_k.weight", h, n_kv * hd) + vname = f"blk.{il}.attn_v.weight" + v_raw = self._gemm_f32(vname, h, n_kv * hd) if vname in self.gg.tensors else k_raw + + q = rms_norm(q.reshape(N, cfg.n_heads, hd), cfg.eps, + w.f32(f"blk.{il}.attn_q_norm.weight")) + k = rms_norm(k_raw.reshape(N, n_kv, hd), cfg.eps, + w.f32(f"blk.{il}.attn_k_norm.weight")) + v = rms_norm(v_raw.reshape(N, n_kv, hd), cfg.eps) + q = rope_neox(q, pos, base, ff) + k = rope_neox(k, pos, base, ff) + dmp("q_pos", il, q); dmp("k_pos", il, k); dmp("v_normed", il, v) + + mask = build_mask(P, C, swa, cfg.window, n_cols=_pad32(N)) + o = self._attention(il, q, k, v, mask) + attn = self._gemm_f32(f"blk.{il}.attn_output.weight", o, cfg.d_model) + attn = rms_norm(attn, cfg.eps, w.f32(f"blk.{il}.post_attention_norm.weight")) + attn_out = (attn + x).astype(F32) + dmp("attn_out", il, attn_out) + + m = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.ffn_norm.weight")) + g_ = gelu_tanh(self._gemm_f32(f"blk.{il}.ffn_gate.weight", m, cfg.n_ff)) + u_ = self._gemm_f32(f"blk.{il}.ffn_up.weight", m, cfg.n_ff) + mlp = self._gemm_f32(f"blk.{il}.ffn_down.weight", g_ * u_, cfg.d_model) + mlp = rms_norm(mlp, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_1.weight")) + dmp("ffn_mlp", il, mlp) + + e_in = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.pre_ffw_norm_2.weight")) + moe = self._moe(il, attn_out, e_in) + moe = rms_norm(moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_2.weight")) + dmp("ffn_moe", il, moe) + + f = rms_norm(mlp + moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm.weight")) + cur = (f + attn_out).astype(F32) + cur[:P] *= w.f32(f"blk.{il}.enc_layer_output_scale.weight")[0] + cur[P:] *= w.f32(f"blk.{il}.layer_output_scale.weight")[0] + dmp("l_out", il, cur) + x = cur + + x = rms_norm(x, cfg.eps, w.f32("output_norm.weight")) + dmp("result_norm", -1, x) + + logits = self._gemm_f32("token_embd.weight", x[P:], cfg.vocab_size) + cap = F32(cfg.final_logit_softcap) + logits = (np.tanh(logits / cap) * cap).astype(F32) + dmp("result_output", -1, logits) + return logits diff --git a/SuperKittens/models/gemma/diffusion/gguf_io.py b/SuperKittens/models/gemma/diffusion/gguf_io.py new file mode 100644 index 0000000..856a054 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/gguf_io.py @@ -0,0 +1,231 @@ +# pyright: reportMissingImports=false +"""gguf_io.py — minimal, dependency-free GGUF reader + numpy K-quant dequant. + +DiffusionGemma is GGUF-canonical (Q4_K_M, 15.65 GiB) and cannot take the +gemma4 dequant-to-bf16 load path on a 16 GB host, so the family reads tensors +straight off the mmap'd file (native quant blocks feed the SK quant GEMM +kernels; dequant here is only for F32 sidecars, embedding rows, and the CPU +reference forward). Dequant formulas mirror ggml's dequantize_row_* exactly. +""" +from __future__ import annotations + +import mmap +import struct +from dataclasses import dataclass + +import numpy as np + +GGUF_MAGIC = 0x46554747 + +# ggml type ids -> (name, block_elems, block_bytes) +GGML_TYPES = { + 0: ("F32", 1, 4), + 1: ("F16", 1, 2), + 2: ("Q4_0", 32, 18), + 8: ("Q8_0", 32, 34), + 12: ("Q4_K", 256, 144), + 13: ("Q5_K", 256, 176), + 14: ("Q6_K", 256, 210), + 6: ("Q5_0", 32, 22), + 30: ("BF16", 1, 2), +} + +_KV_FMT = {0: " str: + return GGML_TYPES[self.ggml_type][0] + + @property + def nbytes(self) -> int: + _, be, bb = GGML_TYPES[self.ggml_type] + n = 1 + for d in self.shape: + n *= d + assert n % be == 0, (self.name, self.shape, self.type_name) + return n // be * bb + + @property + def row_bytes(self) -> int: + """bytes per ne0-row (the K dimension all quant kernels stride by).""" + _, be, bb = GGML_TYPES[self.ggml_type] + assert self.shape[0] % be == 0 + return self.shape[0] // be * bb + + +class GGUFFile: + """mmap-backed GGUF: metadata dict + tensor table + raw/dequant access.""" + + def __init__(self, path: str): + self.path = path + self._f = open(path, "rb") + self.mm = mmap.mmap(self._f.fileno(), 0, prot=mmap.PROT_READ) + self.meta: dict[str, object] = {} + self.tensors: dict[str, TensorInfo] = {} + self._parse() + + # -- parsing ------------------------------------------------------------ + + def _read(self, fmt: str) -> object: + v = struct.unpack_from(fmt, self.mm, self._pos)[0] + self._pos += struct.calcsize(fmt) + return v + + def _read_str(self) -> str: + n = self._read(" None: + self._pos = 0 + magic = self._read("= 2, f"gguf v{version} unsupported" + n_tensors = self._read(" memoryview: + ti = self.tensors[name] + return memoryview(self.mm)[ti.offset:ti.offset + ti.nbytes] + + def dequant(self, name: str, rows: slice | np.ndarray | None = None) -> np.ndarray: + """Dequantize to f32, shaped [shape[::-1]] (numpy row-major: last ggml + dim first). `rows` selects along the OUTER (row) dimension, where a row + is one ne0-slice — exactly ggml's get_rows granularity.""" + ti = self.tensors[name] + k = ti.shape[0] + n_rows = 1 + for d in ti.shape[1:]: + n_rows *= d + rb = ti.row_bytes + buf = np.frombuffer(self.mm, dtype=np.uint8, count=n_rows * rb, + offset=ti.offset).reshape(n_rows, rb) + if rows is not None: + buf = buf[rows] + out = dequant_rows(buf, ti.type_name, k) + if rows is None and len(ti.shape) > 1: + out = out.reshape(tuple(ti.shape[::-1])) + return out + + +# -- numpy dequant (mirrors ggml dequantize_row_*) --------------------------- + +def _f16(u16: np.ndarray) -> np.ndarray: + return u16.view(np.float16).astype(np.float32) + + +def dequant_rows(buf: np.ndarray, type_name: str, k: int) -> np.ndarray: + """buf: uint8 [R, row_bytes] -> f32 [R, k].""" + r = buf.shape[0] + if type_name == "F32": + return buf.reshape(-1).view(np.float32).reshape(r, k).copy() + if type_name == "F16": + return buf.reshape(-1).view(np.float16).astype(np.float32).reshape(r, k) + if type_name == "Q8_0": + nb = k // 32 + b = buf.reshape(r * nb, 34) + d = _f16(b[:, :2].copy().view(np.uint16)[:, 0]) + q = b[:, 2:].view(np.int8).astype(np.float32) + return (q * d[:, None]).reshape(r, k) + if type_name == "Q4_K": + return _dequant_q4k(buf, k) + if type_name == "Q6_K": + return _dequant_q6k(buf, k) + raise ValueError(f"dequant for {type_name} not implemented") + + +def _dequant_q4k(buf: np.ndarray, k: int) -> np.ndarray: + r = buf.shape[0] + nb = k // 256 + b = buf.reshape(r * nb, 144) + d = _f16(b[:, 0:2].copy().view(np.uint16)[:, 0]) # [B] + dmin = _f16(b[:, 2:4].copy().view(np.uint16)[:, 0]) + sc_raw = b[:, 4:16].astype(np.uint16) # [B, 12] + qs = b[:, 16:144] # [B, 128] + # get_scale_min_k4 for j = 0..7 + sc = np.empty((b.shape[0], 8), np.float32) + mn = np.empty((b.shape[0], 8), np.float32) + for j in range(4): + sc[:, j] = (sc_raw[:, j] & 63).astype(np.float32) + mn[:, j] = (sc_raw[:, j + 4] & 63).astype(np.float32) + for j in range(4, 8): + sc[:, j] = ((sc_raw[:, j + 4] & 0x0F) | ((sc_raw[:, j - 4] >> 6) << 4)).astype(np.float32) + mn[:, j] = ((sc_raw[:, j + 4] >> 4) | ((sc_raw[:, j] >> 6) << 4)).astype(np.float32) + lo = (qs & 0x0F).astype(np.float32).reshape(-1, 4, 32) # byte group g -> sub-block 2g + hi = (qs >> 4).astype(np.float32).reshape(-1, 4, 32) # -> sub-block 2g+1 + y = np.empty((b.shape[0], 8, 32), np.float32) + y[:, 0::2, :] = lo + y[:, 1::2, :] = hi + y = y * (d[:, None] * sc)[:, :, None] - (dmin[:, None] * mn)[:, :, None] + return y.reshape(r, k) + + +def _dequant_q6k(buf: np.ndarray, k: int) -> np.ndarray: + r = buf.shape[0] + nb = k // 256 + b = buf.reshape(r * nb, 210) + ql = b[:, 0:128] + qh = b[:, 128:192] + scales = b[:, 192:208].view(np.int8).astype(np.float32) # [B, 16] + d = _f16(b[:, 208:210].copy().view(np.uint16)[:, 0]) + y = np.empty((b.shape[0], 256), np.float32) + for half in range(2): + qlh = ql[:, 64 * half:64 * half + 64] + qhh = qh[:, 32 * half:32 * half + 32] + sch = scales[:, 8 * half:8 * half + 8] + l = np.arange(32) + is_ = l >> 4 # [32] in {0,1} + q1 = ((qlh[:, :32] & 0x0F) | (((qhh >> 0) & 3) << 4)).astype(np.int32) - 32 + q2 = ((qlh[:, 32:] & 0x0F) | (((qhh >> 2) & 3) << 4)).astype(np.int32) - 32 + q3 = ((qlh[:, :32] >> 4) | (((qhh >> 4) & 3) << 4)).astype(np.int32) - 32 + q4 = ((qlh[:, 32:] >> 4) | (((qhh >> 6) & 3) << 4)).astype(np.int32) - 32 + base = 128 * half + y[:, base + 0:base + 32] = sch[:, is_ + 0] * q1 + y[:, base + 32:base + 64] = sch[:, is_ + 2] * q2 + y[:, base + 64:base + 96] = sch[:, is_ + 4] * q3 + y[:, base + 96:base + 128] = sch[:, is_ + 6] * q4 + y *= d[:, None] + return y.reshape(r, k) diff --git a/SuperKittens/models/gemma/diffusion/graph_ref.py b/SuperKittens/models/gemma/diffusion/graph_ref.py new file mode 100644 index 0000000..1b6748e --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/graph_ref.py @@ -0,0 +1,235 @@ +# pyright: reportMissingImports=false +"""graph_ref.py — host math shared by the CPU reference forward and the Metal +driver's glue, plus the full CPU-f32 reference forward. + +Every op mirrors llama.cpp PR #24423 (src/models/diffusion-gemma.cpp + +gemma4-common.h + llm_graph_context::build_moe_ffn at c84e85af). The CPU +forward is the parity bisect oracle: dump(name, il, arr) taps match the GPU +driver's taps one-for-one, so first-divergence is a numpy diff per layer. +""" +from __future__ import annotations + +import numpy as np + +from .config import DiffusionGemmaConfig +from .gguf_io import GGUFFile + +F32 = np.float32 + + +# -- primitive ops (all f32 in / f32 out) ------------------------------------- + +def rms_norm(x: np.ndarray, eps: float, w: np.ndarray | None = None) -> np.ndarray: + inv = 1.0 / np.sqrt((x.astype(F32) ** 2).mean(axis=-1, keepdims=True) + F32(eps)) + y = x * inv + if w is not None: + y = y * w + return y.astype(F32) + + +def gelu_tanh(x: np.ndarray) -> np.ndarray: + # ggml_gelu / HF gelu_pytorch_tanh + x = x.astype(F32) + return (0.5 * x * (1.0 + np.tanh(0.7978845608028654 * (x + 0.044715 * x ** 3)))).astype(F32) + + +def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray: + m = x.max(axis=axis, keepdims=True) + e = np.exp((x - m).astype(F32)) + return (e / e.sum(axis=axis, keepdims=True)).astype(F32) + + +def rope_neox(x: np.ndarray, pos: np.ndarray, base: float, + freq_factors: np.ndarray | None) -> np.ndarray: + """x [T, H, D] -> rotated. ggml NEOX: pair (j, j+D/2), angle = + pos * base^(-2j/D) / ff[j]. Full-dim rotation (n_rot == head_dim here).""" + T, H, D = x.shape + j = np.arange(D // 2, dtype=np.float64) + inv = np.power(float(base), -2.0 * j / D) + if freq_factors is not None: + inv = inv / freq_factors.astype(np.float64) + ang = pos.astype(np.float64)[:, None] * inv[None, :] # [T, D/2] + c = np.cos(ang).astype(F32)[:, None, :] + s = np.sin(ang).astype(F32)[:, None, :] + x0 = x[..., :D // 2] + x1 = x[..., D // 2:] + out = np.empty_like(x, dtype=F32) + out[..., :D // 2] = x0 * c - x1 * s + out[..., D // 2:] = x0 * s + x1 * c + return out + + +def build_mask(P: int, C: int, swa: bool, n_swa: int, n_cols: int | None = None) -> np.ndarray: + """Additive mask [P+C, n_cols] f32 (0 / -inf), PR rules. Prompt queries: + causal over prompt only (SWA-clipped when swa). Canvas queries: + bidirectional — global sees all; SWA sees last (n_swa-1) prompt + canvas. + Columns >= P+C (pad) stay -inf.""" + N = P + C + n_cols = n_cols or N + mask = np.full((N, n_cols), -np.inf, dtype=F32) + q = np.arange(N)[:, None] + k = np.arange(N)[None, :] + q_canvas = q >= P + k_canvas = k >= P + canvas_prompt_lo = P - n_swa + 1 + if swa: + allow_canvas_q = k_canvas | (k >= canvas_prompt_lo) + else: + allow_canvas_q = np.ones((1, N), dtype=bool) + allow_prompt_q = (~k_canvas) & (k <= q) + if swa: + allow_prompt_q = allow_prompt_q & (q - k < n_swa) # is_masked_swa STANDARD + allow = np.where(q_canvas, allow_canvas_q, allow_prompt_q) + mask[:, :N][allow] = 0.0 + return mask + + +# -- weight access ------------------------------------------------------------- + +class Weights: + """Thin named access over the GGUF (dequant cached only for small F32).""" + + def __init__(self, gg: GGUFFile): + self.gg = gg + self._f32: dict[str, np.ndarray] = {} + + def f32(self, name: str) -> np.ndarray: + a = self._f32.get(name) + if a is None: + ti = self.gg.tensors[name] + assert ti.type_name == "F32", name + a = self.gg.dequant(name) + self._f32[name] = a + return a + + def dq(self, name: str, rows=None) -> np.ndarray: + return self.gg.dequant(name, rows) + + +# -- region-aware unified forward (CPU f32 reference) ------------------------- + +def embed_tokens(w: Weights, cfg: DiffusionGemmaConfig, ids: np.ndarray, P: int) -> np.ndarray: + x = w.dq("token_embd.weight", rows=np.asarray(ids, np.int64)) + x = x * F32(np.sqrt(F32(cfg.d_model))) + x[P:] = rms_norm(x[P:], cfg.eps) # canvas rows: rmsnorm no-scale (zero-SC) + return x.astype(F32) + + +def moe_route(w: Weights, cfg: DiffusionGemmaConfig, attn_out: np.ndarray, il: int): + """Router (operates on the UNNORMED residual): rms_noscale -> /sqrt(d) -> + * gate_inp scale -> logits -> softmax -> top-8 -> renorm weights.""" + t = rms_norm(attn_out, cfg.eps) + t = t * F32(1.0 / np.sqrt(F32(cfg.d_model))) + t = t * w.f32(f"blk.{il}.ffn_gate_inp.scale") + logits = t @ w.f32(f"blk.{il}.ffn_gate_inp.weight").T # [T, 128] + probs = softmax(logits) + sel = np.argsort(-probs, axis=-1, kind="stable")[:, :cfg.n_expert_used] # [T, 8] + wts = np.take_along_axis(probs, sel, axis=-1) + wts = wts / np.maximum(wts.sum(-1, keepdims=True), F32(6.103515625e-5)) + return sel, wts.astype(F32) + + +def forward_cpu(gg: GGUFFile, cfg: DiffusionGemmaConfig, ids: np.ndarray, + P: int, dump=None) -> np.ndarray: + """Unified [prompt|canvas] zero-SC forward; returns canvas logits f32 [C, V].""" + w = Weights(gg) + N = len(ids) + C = N - P + pos = np.arange(N, dtype=np.int64) + dmp = dump or (lambda name, il, arr: None) + + x = embed_tokens(w, cfg, np.asarray(ids), P) + dmp("inp_region", -1, x) + + rope_ff = w.f32("rope_freqs.weight") + + for il in range(cfg.n_layers): + swa = cfg.is_swa[il] + hd = cfg.head_dim(il) + n_kv = cfg.n_kv_heads[il] + base, use_ff = cfg.rope_params(il) + ff = rope_ff if use_ff else None + gqa = cfg.n_heads // n_kv + + h = rms_norm(x, cfg.eps, w.f32(f"blk.{il}.attn_norm.weight")) + + q = (h @ w.dq(f"blk.{il}.attn_q.weight").T).reshape(N, cfg.n_heads, hd) + k_raw = h @ w.dq(f"blk.{il}.attn_k.weight").T + vname = f"blk.{il}.attn_v.weight" + v_raw = (h @ w.dq(vname).T) if vname in gg.tensors else k_raw # global: V = raw k_proj + k = k_raw.reshape(N, n_kv, hd) + v = v_raw.reshape(N, n_kv, hd) + + q = rms_norm(q, cfg.eps, w.f32(f"blk.{il}.attn_q_norm.weight")) + k = rms_norm(k, cfg.eps, w.f32(f"blk.{il}.attn_k_norm.weight")) + v = rms_norm(v, cfg.eps) # v-norm: no scale + q = rope_neox(q, pos, base, ff) + k = rope_neox(k, pos, base, ff) + dmp("q_pos", il, q); dmp("k_pos", il, k); dmp("v_normed", il, v) + + mask = build_mask(P, C, swa, cfg.window) + o = np.empty((N, cfg.n_heads, hd), dtype=F32) + for hh in range(cfg.n_heads): + g = hh // gqa + s = (q[:, hh, :] @ k[:, g, :].T) * F32(cfg.attn_scale) + mask + a = softmax(s) + o[:, hh, :] = a @ v[:, g, :] + attn = o.reshape(N, cfg.n_heads * hd) @ w.dq(f"blk.{il}.attn_output.weight").T + attn = rms_norm(attn, cfg.eps, w.f32(f"blk.{il}.post_attention_norm.weight")) + attn_out = (attn + x).astype(F32) + dmp("attn_out", il, attn_out) + + # dense MLP (shared expert): geglu, then post_ffw_norm_1 + m = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.ffn_norm.weight")) + g_ = gelu_tanh(m @ w.dq(f"blk.{il}.ffn_gate.weight").T) + u_ = m @ w.dq(f"blk.{il}.ffn_up.weight").T + mlp = (g_ * u_) @ w.dq(f"blk.{il}.ffn_down.weight").T + mlp = rms_norm(mlp, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_1.weight")) + dmp("ffn_mlp", il, mlp) + + # MoE + e_in = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.pre_ffw_norm_2.weight")) + sel, wts = moe_route(w, cfg, attn_out, il) + dmp("moe_sel", il, sel.astype(np.int32)); dmp("moe_wts", il, wts) + down_s = w.f32(f"blk.{il}.ffn_down_exps.scale") + moe = np.zeros((N, cfg.d_model), dtype=F32) + gu_ti = gg.tensors[f"blk.{il}.ffn_gate_up_exps.weight"] + dn_ti = gg.tensors[f"blk.{il}.ffn_down_exps.weight"] + for e in np.unique(sel): + tok, slot = np.nonzero(sel == e) + r0 = e * 2 * cfg.n_ff_exp + wgu = w.dq(gu_ti.name, rows=slice(r0, r0 + 2 * cfg.n_ff_exp)) + gu = e_in[tok] @ wgu.T # [m, 1408] + act = gelu_tanh(gu[:, :cfg.n_ff_exp]) * gu[:, cfg.n_ff_exp:] + r0 = e * cfg.d_model + wdn = w.dq(dn_ti.name, rows=slice(r0, r0 + cfg.d_model)) + d_ = (act @ wdn.T) * down_s[e] # [m, 2816] + moe[tok] += d_ * wts[tok, slot][:, None] + moe = rms_norm(moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_2.weight")) + dmp("ffn_moe", il, moe) + + f = rms_norm(mlp + moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm.weight")) + cur = (f + attn_out).astype(F32) + + enc_s = w.f32(f"blk.{il}.enc_layer_output_scale.weight")[0] + dec_s = w.f32(f"blk.{il}.layer_output_scale.weight")[0] + cur[:P] *= enc_s + cur[P:] *= dec_s + dmp("l_out", il, cur) + x = cur + + x = rms_norm(x, cfg.eps, w.f32("output_norm.weight")) + dmp("result_norm", -1, x) + + # tied head on canvas rows only, vocab-chunked (full f32 dequant is ~3 GB) + xc = x[P:] + V = cfg.vocab_size + logits = np.empty((C, V), dtype=F32) + step = 16384 + for v0 in range(0, V, step): + wchunk = w.dq("token_embd.weight", rows=slice(v0, min(v0 + step, V))) + logits[:, v0:v0 + wchunk.shape[0]] = xc @ wchunk.T + cap = F32(cfg.final_logit_softcap) + logits = (np.tanh(logits / cap) * cap).astype(F32) + dmp("result_output", -1, logits) + return logits diff --git a/SuperKittens/models/gemma/diffusion/runner.py b/SuperKittens/models/gemma/diffusion/runner.py new file mode 100644 index 0000000..0a4a7bf --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/runner.py @@ -0,0 +1,84 @@ +# pyright: reportMissingImports=false +"""runner.py — Stage-1 parity CLI for DiffusionGemma. + +Feeds golden token ids [prompt | canvas] through one unified zero-SC forward +and writes the canvas logits as raw f32 — byte-compatible with the llama.cpp +PR #24423 `llama-diffusion-gemma-eval` harness, so the two outputs diff +directly. + + python -m SuperKittens.models.gemma.diffusion.runner \ + --gguf model.gguf --prompt-ids p.i32 --canvas-ids c.i32 \ + --out logits.bin [--mode gpu|cpu] [--dump-dir d --dump-names l_out,...] +""" +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +from .config import config_from_gguf +from .gguf_io import GGUFFile + + +def make_dump(dump_dir: str | None, names: set[str]): + if not dump_dir: + return None + d = Path(dump_dir) + d.mkdir(parents=True, exist_ok=True) + + def dump(name: str, il: int, arr: np.ndarray): + if names and name not in names: + return + np.save(d / f"{name}.{il}.npy", np.asarray(arr)) + return dump + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--gguf", required=True) + ap.add_argument("--prompt-ids", required=True) + ap.add_argument("--canvas-ids", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--mode", choices=["gpu", "cpu"], default="gpu") + ap.add_argument("--layers", type=int, default=0, help="truncate to first N layers (debug)") + ap.add_argument("--dump-dir", default=None) + ap.add_argument("--dump-names", default="l_out") + args = ap.parse_args(argv) + + prompt = np.fromfile(args.prompt_ids, dtype=np.int32) + canvas = np.fromfile(args.canvas_ids, dtype=np.int32) + ids = np.concatenate([prompt, canvas]) + P = len(prompt) + + gg = GGUFFile(args.gguf) + cfg = config_from_gguf(gg.meta) + if len(canvas) != cfg.canvas_length: + print(f"canvas len {len(canvas)} != model canvas_length {cfg.canvas_length}", + file=sys.stderr) + return 1 + if args.layers: + cfg.n_layers = args.layers + dump = make_dump(args.dump_dir, set(filter(None, args.dump_names.split(",")))) + + t0 = time.time() + if args.mode == "cpu": + from .graph_ref import forward_cpu + logits = forward_cpu(gg, cfg, ids, P, dump=dump) + else: + from .forward_metal import DiffusionGemmaMetal + m = DiffusionGemmaMetal(args.gguf, cfg) + m.dump = dump + logits = m.forward(ids, P) + dt = time.time() - t0 + + logits.astype(np.float32).tofile(args.out) + print(f"wrote {logits.shape[0]} x {logits.shape[1]} f32 logits to {args.out} " + f"({args.mode}, P={P}, {dt:.1f}s)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/temp/diffgemma_s1/test_ops.py b/temp/diffgemma_s1/test_ops.py new file mode 100644 index 0000000..c25918e --- /dev/null +++ b/temp/diffgemma_s1/test_ops.py @@ -0,0 +1,133 @@ +# pyright: reportMissingImports=false +"""Synthetic op tests for the DiffusionGemma Stage-1 driver. + +Cross-validates: gemm_mma_{f16,q8_0,q4k,q6k} dispatch plumbing vs the numpy +dequant in gguf_io (random quant blocks — both sides must agree), the +dg_softmax_mask kernel vs numpy, and the GEMM-composed attention vs a pure +numpy reference. No model weights needed; runs on any Apple Silicon box. +""" +import sys +from pathlib import Path + +import numpy as np + +SK_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(SK_ROOT)) + +from SuperKittens.models.gemma.diffusion.gguf_io import dequant_rows # noqa: E402 +from SuperKittens.models.gemma.diffusion.graph_ref import build_mask, softmax # noqa: E402 +from SuperKittens.models.gemma.diffusion import forward_metal as fm # noqa: E402 + +rng = np.random.default_rng(0) +ctx = fm.MetalCtx() + + +def gemm(kernel, a_f16, w_buf_arr, M, N, K, ldc=None): + a_buf = ctx.buf_from(a_f16) + w_buf = ctx.buf_from(w_buf_arr) + c_buf = ctx.buf_empty(M * (ldc or N) * 2) + b = fm.GemmBatch(ctx) + b.gemm(kernel, a_buf, 0, w_buf, 0, c_buf, 0, M, N, K, ldc) + b.run() + return ctx.read(c_buf, np.float16, (M, ldc or N)) + + +def check(name, got, want, tol): + got = got.astype(np.float64) + want = want.astype(np.float64) + rms_rel = np.linalg.norm(got - want) / max(np.linalg.norm(want), 1e-12) + max_abs = np.abs(got - want).max() + status = "OK " if rms_rel < tol else "FAIL" + print(f"{status} {name}: rms rel {rms_rel:.3e} max abs {max_abs:.3e} (tol {tol})") + return rms_rel < tol + + +def rand_q8_rows(n, k): + nb = k // 32 + raw = np.zeros((n, nb, 34), np.uint8) + d = (rng.uniform(0.001, 0.02, (n, nb)).astype(np.float16)) + raw[:, :, :2] = d.view(np.uint8).reshape(n, nb, 2) + raw[:, :, 2:] = rng.integers(0, 256, (n, nb, 32), dtype=np.uint8) + return raw.reshape(n, nb * 34) + + +def rand_q4k_rows(n, k): + nb = k // 256 + raw = np.zeros((n, nb, 144), np.uint8) + d = rng.uniform(0.001, 0.02, (n, nb)).astype(np.float16) + dmin = rng.uniform(0.0005, 0.01, (n, nb)).astype(np.float16) + raw[:, :, 0:2] = d.view(np.uint8).reshape(n, nb, 2) + raw[:, :, 2:4] = dmin.view(np.uint8).reshape(n, nb, 2) + raw[:, :, 4:] = rng.integers(0, 256, (n, nb, 140), dtype=np.uint8) + return raw.reshape(n, nb * 144) + + +def rand_q6k_rows(n, k): + nb = k // 256 + raw = rng.integers(0, 256, (n, nb, 210), dtype=np.uint8) + d = rng.uniform(0.0005, 0.005, (n, nb)).astype(np.float16) + raw[:, :, 208:210] = d.view(np.uint8).reshape(n, nb, 2) + return raw.reshape(n, nb * 210) + + +ok = True + +# f16 GEMM, ragged M/N + ldc band +M, K, N = 37, 512, 96 +a = rng.standard_normal((M, K)).astype(np.float16) +w = rng.standard_normal((N, K)).astype(np.float16) +got = gemm("gemm_mma_f16", a, w, M, N, K) +want = a.astype(np.float32) @ w.astype(np.float32).T +ok &= check("gemm_mma_f16", got, want, 2e-2) + +# quant GEMMs vs numpy dequant (mutual validation of kernel + gguf_io) +for tname, kern, gen in [("Q8_0", "gemm_mma_q8_0", rand_q8_rows), + ("Q4_K", "gemm_mma_q4k", rand_q4k_rows), + ("Q6_K", "gemm_mma_q6k", rand_q6k_rows)]: + K = 768 if tname != "Q8_0" else 704 + N, M = 95, 33 + raw = gen(N, K) + wf = dequant_rows(raw, tname, K) + a = (rng.standard_normal((M, K)) * 0.1).astype(np.float16) + got = gemm(kern, a, raw.reshape(-1), M, N, K) + want = a.astype(np.float32) @ wf.T + ok &= check(f"gemm_mma {tname}", got, want, 2e-2) + +# masked softmax kernel vs numpy (incl pad cols) +Ntok, Np, H = 69, 96, 4 +s = (rng.standard_normal((H * Ntok, Np)) * 4).astype(np.float16) +mask = build_mask(13, Ntok - 13, True, 24, n_cols=Np) +s_buf = ctx.buf_from(s) +m_buf = ctx.buf_from(mask) +b = fm.GemmBatch(ctx) +b.softmax_mask(s_buf, m_buf, rows=H * Ntok, ncols=Np, ntok=Ntok, scale=1.0) +b.run() +got = ctx.read(s_buf, np.float16, (H * Ntok, Np)) +want = softmax(np.tile(mask, (H, 1)) + s.astype(np.float32)) +ok &= check("dg_softmax_mask", got, want, 2e-2) + +# GEMM-composed attention vs numpy (GQA, dual dims, region mask) +class FakeCfg: + attn_scale = 1.0 + +for hd, n_kv in [(256, 8), (512, 2)]: + P, C = 11, 53 + N = P + C + Hq = 16 + q = rng.standard_normal((N, Hq, hd)).astype(np.float32) * 0.3 + k = rng.standard_normal((N, n_kv, hd)).astype(np.float32) * 0.3 + v = rng.standard_normal((N, n_kv, hd)).astype(np.float32) * 0.3 + mask = build_mask(P, C, n_kv == 8, 1024, n_cols=fm._pad32(N)) + drv = object.__new__(fm.DiffusionGemmaMetal) + drv.cfg = FakeCfg() + drv.ctx = ctx + got = drv._attention(0, q, k, v, mask) + want = np.empty((N, Hq, hd), np.float32) + for h in range(Hq): + g = h // (Hq // n_kv) + sc = q[:, h] @ k[:, g].T + mask[:, :N] + want[:, h] = softmax(sc) @ v[:, g] + ok &= check(f"attention hd={hd} kv={n_kv}", got, want.reshape(N, Hq * hd), 3e-2) + +print("ALL OK" if ok else "FAILURES", flush=True) +sys.exit(0 if ok else 1) From f0f6fc37ac6679519032e2301c2356f2c1fdd0a0 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 18:34:58 -0400 Subject: [PATCH 09/31] diffgemma stage1: 1-d tensor dequant shape fix; parity/bisect tooling (compare_logits, compare_dumps, run_sk) --- .../models/gemma/diffusion/gguf_io.py | 4 +- temp/diffgemma_s1/compare_dumps.py | 28 +++++++++++ temp/diffgemma_s1/compare_logits.py | 46 +++++++++++++++++++ temp/diffgemma_s1/run_sk.sh | 22 +++++++++ 4 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 temp/diffgemma_s1/compare_dumps.py create mode 100644 temp/diffgemma_s1/compare_logits.py create mode 100644 temp/diffgemma_s1/run_sk.sh diff --git a/SuperKittens/models/gemma/diffusion/gguf_io.py b/SuperKittens/models/gemma/diffusion/gguf_io.py index 856a054..149b91b 100644 --- a/SuperKittens/models/gemma/diffusion/gguf_io.py +++ b/SuperKittens/models/gemma/diffusion/gguf_io.py @@ -146,8 +146,8 @@ def dequant(self, name: str, rows: slice | np.ndarray | None = None) -> np.ndarr if rows is not None: buf = buf[rows] out = dequant_rows(buf, ti.type_name, k) - if rows is None and len(ti.shape) > 1: - out = out.reshape(tuple(ti.shape[::-1])) + if rows is None: + out = out.reshape(tuple(ti.shape[::-1])) # 1-d tensors flatten to [k] return out diff --git a/temp/diffgemma_s1/compare_dumps.py b/temp/diffgemma_s1/compare_dumps.py new file mode 100644 index 0000000..3bb0da9 --- /dev/null +++ b/temp/diffgemma_s1/compare_dumps.py @@ -0,0 +1,28 @@ +# pyright: reportMissingImports=false +"""compare_dumps.py — first-divergence bisect between two dump dirs +(runner --dump-dir for gpu vs cpu modes). Prints rms-rel per tap in layer +order so the first bad layer/op is obvious. + + python3 compare_dumps.py dir_a dir_b +""" +import sys +from pathlib import Path + +import numpy as np + +a_dir, b_dir = Path(sys.argv[1]), Path(sys.argv[2]) +names = sorted(set(p.name for p in a_dir.glob("*.npy")) & set(p.name for p in b_dir.glob("*.npy")), + key=lambda n: (int(n.split(".")[-2]), n)) +for n in names: + a = np.load(a_dir / n).astype(np.float64) + b = np.load(b_dir / n).astype(np.float64) + if a.shape != b.shape: + print(f"{n}: SHAPE {a.shape} vs {b.shape}") + continue + if a.dtype.kind == "i" or n.startswith("moe_sel"): + mism = (a != b).mean() + print(f"{n}: sel mismatch {100*mism:.3f}%") + continue + rms = np.linalg.norm(a - b) / max(np.linalg.norm(b), 1e-12) + mx = np.abs(a - b).max() + print(f"{n}: rms_rel {rms:.3e} max_abs {mx:.3e}") diff --git a/temp/diffgemma_s1/compare_logits.py b/temp/diffgemma_s1/compare_logits.py new file mode 100644 index 0000000..adb35db --- /dev/null +++ b/temp/diffgemma_s1/compare_logits.py @@ -0,0 +1,46 @@ +# pyright: reportMissingImports=false +"""compare_logits.py — SK vs llama.cpp canvas-logit parity report. + +Both files are raw f32 [C, n_vocab] (canvas rows). Reports per-position rel +err stats + argmax identity (the meaningful bar at Q4) + top-5 overlap. + + python3 compare_logits.py ref.bin sk.bin [n_vocab] +""" +import sys + +import numpy as np + +V = int(sys.argv[3]) if len(sys.argv) > 3 else 262144 +ref = np.fromfile(sys.argv[1], np.float32).reshape(-1, V) +got = np.fromfile(sys.argv[2], np.float32).reshape(-1, V) +assert ref.shape == got.shape, (ref.shape, got.shape) +C = ref.shape[0] + +diff = np.abs(got - ref) +denom = np.maximum(np.abs(ref), 1e-3) +rel = diff / denom +rel_pos_max = rel.max(axis=1) +rel_pos_mean = rel.mean(axis=1) +rms_rel = np.linalg.norm(got - ref, axis=1) / np.linalg.norm(ref, axis=1) + +am_ref = ref.argmax(axis=1) +am_got = got.argmax(axis=1) +am_match = (am_ref == am_got) + +top5_ref = np.argsort(-ref, axis=1)[:, :5] +top5_got = np.argsort(-got, axis=1)[:, :5] +t5 = np.array([len(np.intersect1d(top5_ref[i], top5_got[i])) for i in range(C)]) + +print(f"positions : {C}") +print(f"rel err max/pos : mean {rel_pos_max.mean():.4f} median {np.median(rel_pos_max):.4f} worst {rel_pos_max.max():.4f}") +print(f"rel err mean/pos : mean {rel_pos_mean.mean():.5f} worst {rel_pos_mean.max():.5f}") +print(f"rms rel per pos : mean {rms_rel.mean():.5f} worst {rms_rel.max():.5f}") +print(f"max abs diff : {diff.max():.4f} (logit scale, softcap 30)") +print(f"ARGMAX match : {am_match.sum()}/{C} = {100.0*am_match.mean():.2f}%") +print(f"top5 overlap : mean {t5.mean():.2f}/5 min {t5.min()}") +if not am_match.all(): + bad = np.nonzero(~am_match)[0][:10] + for i in bad: + print(f" pos {i}: ref argmax {am_ref[i]} ({ref[i, am_ref[i]]:.3f}) vs " + f"got {am_got[i]} ({got[i, am_got[i]]:.3f}); " + f"got[ref_am]={got[i, am_ref[i]]:.3f} ref[got_am]={ref[i, am_got[i]]:.3f}") diff --git a/temp/diffgemma_s1/run_sk.sh b/temp/diffgemma_s1/run_sk.sh new file mode 100644 index 0000000..4eeb62d --- /dev/null +++ b/temp/diffgemma_s1/run_sk.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# SK GPU forward for the 3 parity prompts on amelia (run detached via caffeinate). +# Usage: run_sk.sh [gpu|cpu] [prompt_index|all] [extra runner args...] +exec >> ~/sk-diffg-s1/sk_run.log 2>&1 +set -x +MODE=${1:-gpu} +WHICH=${2:-all} +shift 2 || true +cd ~/sk-diffg-s1 +export PYTHONPATH=~/sk-diffg-s1 +for i in 1 2 3; do + if [ "$WHICH" != "all" ] && [ "$WHICH" != "$i" ]; then continue; fi + date + sysctl vm.swapusage + /usr/bin/python3 -m SuperKittens.models.gemma.diffusion.runner \ + --gguf ~/diffgemma-gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf \ + --prompt-ids inputs/p${i}_prompt.i32 --canvas-ids inputs/p${i}_canvas.i32 \ + --out sk_${MODE}_p${i}.bin --mode ${MODE} "$@" + echo "SK_${MODE}_P${i}_RC=$?" +done +date +echo SK_RUN_DONE From f1381072cf0157ba570d0cb9afe8191c232c0c6b Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 18:43:43 -0400 Subject: [PATCH 10/31] diffgemma stage1: Q5_0 support (16/30 layers in this Q4_K_M mix), f32-out QK^T + f32-in softmax (kq range parity with ggml PREC_F32), EOF-safe mmap windows --- .../models/gemma/diffusion/dg_kernels.metal | 152 +++++++++++++++--- .../models/gemma/diffusion/forward_metal.py | 40 +++-- .../models/gemma/diffusion/gguf_io.py | 18 +++ temp/diffgemma_s1/STATUS.md | 85 ++++++++++ temp/diffgemma_s1/compare_logits.py | 10 ++ temp/diffgemma_s1/test_ops.py | 37 ++++- 6 files changed, 299 insertions(+), 43 deletions(-) create mode 100644 temp/diffgemma_s1/STATUS.md diff --git a/SuperKittens/models/gemma/diffusion/dg_kernels.metal b/SuperKittens/models/gemma/diffusion/dg_kernels.metal index ea7f537..08d64e8 100644 --- a/SuperKittens/models/gemma/diffusion/dg_kernels.metal +++ b/SuperKittens/models/gemma/diffusion/dg_kernels.metal @@ -1,24 +1,136 @@ // dg_kernels.metal — DiffusionGemma family kernels. // -// dg_softmax_mask: masked row softmax for the GEMM-composed unified attention -// (QK^T tiles come from kernels/gemm/gemm_mma.metal at M = P+C). The D=128 -// production attention kernels don't apply at head_dim 256/512, so Stage 1 -// runs QK^T -> this kernel -> @V as three plain dispatches. +// Runtime-compiled CONCATENATED AFTER kernels/gemm/gemm_mma.metal (see +// forward_metal.MetalCtx), so the skmma tile loaders and GEMM_MMA_BODY macro +// are in scope — family-only variants live here without touching the shared +// kernel tree. // -// S : fp16 [R, ncols], R = n_heads * n_tok (head-major), softmaxed in place -// mask : f32 [n_tok, ncols] additive (0 / -inf); row r uses mask row r % n_tok -// (one mask per layer, shared across heads; pad cols carry -inf) -// scale: kq scale applied before the mask (1.0 for this family — qk-norm) +// dg_gemm_mma_q5_0 : gemm_mma for Q5_0 weights (this GGUF's Q4_K_M mix puts +// ffn_down / ffn_down_exps at Q5_0 on 16 of 30 layers). +// dg_gemm_qkt_f32 : f16 GEMM with f32 C. QK^T scores need fp32 range — +// ggml forces GGML_PREC_F32 for kq; a half store can +// overflow (no 1/sqrt(d) pre-scale in this family) and +// costs softmax precision. +// dg_softmax_mask : masked row softmax, f32 scores in -> f16 probs out. -#include -using namespace metal; +// ── Q5_0 tile loader. block = 32 weights, 22 B: half d, u32 qh, 16 nibble +// bytes. w = d * (((qs nibble) | (qh bit << 4)) - 16); element wk uses qh +// bit wk in both halves (ggml dequantize_row_q5_0). +struct __attribute__((packed)) dg_q5_0_block { + half d; + uint8_t qh[4]; + uint8_t qs[16]; +}; +inline void load_W_q5_0(threadgroup half* Ws, + device const uchar* W, + uint bc, uint k0, uint N, uint K, + uint lid) +{ + const uint nb = K / 32; + const uint kb = k0 / 32; + for (uint i = lid; i < BN * BK; i += 64) { + const uint wr = i / BK; + const uint wk = i % BK; + const uint gn = bc + wr; + half v = half(0); + if (gn < N) { + device const dg_q5_0_block* blk = + (device const dg_q5_0_block*)(W + ((size_t)gn * nb + kb) * sizeof(dg_q5_0_block)); + const uint qh = (uint)blk->qh[0] | ((uint)blk->qh[1] << 8) + | ((uint)blk->qh[2] << 16) | ((uint)blk->qh[3] << 24); + const uint lo = (wk < 16) ? (blk->qs[wk] & 0x0F) : (blk->qs[wk - 16] >> 4); + const int q = (int)(lo | (((qh >> wk) & 1u) << 4)); + v = (half)((float)blk->d * (float)(q - 16)); + } + Ws[wk * BN + wr] = v; + } +} + +[[host_name("dg_gemm_mma_q5_0")]] +[[kernel]] +void dg_gemm_mma_q5_0( + device const half* A [[buffer(0)]], + device const uchar* W [[buffer(1)]], + device half* C [[buffer(2)]], + constant uint& M [[buffer(3)]], + constant uint& N [[buffer(4)]], + constant uint& K [[buffer(5)]], + constant uint& ldC [[buffer(6)]], + uint2 gid [[threadgroup_position_in_grid]], + uint simd [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) +{ + GEMM_MMA_BODY(load_W_q5_0) +} + +// ── GEMM_MMA_BODY with a float C store (only the final cast differs). +#define DG_GEMM_MMA_BODY_F32OUT(LOAD_W) \ + const uint br = gid.y * BM; \ + const uint bc = gid.x * BN; \ + threadgroup half As[BM * BK]; \ + threadgroup half Ws[BK * BN]; \ + const uint lid = simd * 32 + lane; \ + const uint c0 = simd * MC * 8; \ + simdgroup_float8x8 acc[MR][MC] = {}; \ + for (uint k0 = 0; k0 < K; k0 += BK) { \ + load_A(As, A, br, k0, M, K, K, lid); \ + LOAD_W(Ws, W, bc, k0, N, K, lid); \ + threadgroup_barrier(mem_flags::mem_threadgroup); \ + for (uint k = 0; k < BK / 8; ++k) { \ + simdgroup_half8x8 a[MR]; \ + for (uint r = 0; r < MR; ++r) \ + simdgroup_load(a[r], As + (r * 8) * BK + k * 8, BK); \ + for (uint c = 0; c < MC; ++c) { \ + simdgroup_half8x8 b; \ + simdgroup_load(b, Ws + (k * 8) * BN + c0 + c * 8, BN); \ + for (uint r = 0; r < MR; ++r) \ + simdgroup_multiply_accumulate(acc[r][c], a[r], b, acc[r][c]); \ + } \ + } \ + threadgroup_barrier(mem_flags::mem_threadgroup); \ + } \ + threadgroup float Cs[BM * BN]; \ + for (uint r = 0; r < MR; ++r) \ + for (uint c = 0; c < MC; ++c) \ + simdgroup_store(acc[r][c], Cs + (r * 8) * BN + c0 + c * 8, BN); \ + threadgroup_barrier(mem_flags::mem_threadgroup); \ + for (uint i = lid; i < BM * BN; i += 64) { \ + const uint r = i / BN, cc = i % BN; \ + const uint gr = br + r, gc = bc + cc; \ + if (gr < M && gc < N) C[(size_t)gr * ldC + gc] = Cs[i]; \ + } + +[[host_name("dg_gemm_qkt_f32")]] +[[kernel]] +void dg_gemm_qkt_f32( + device const half* A [[buffer(0)]], + device const half* W [[buffer(1)]], + device float* C [[buffer(2)]], + constant uint& M [[buffer(3)]], + constant uint& N [[buffer(4)]], + constant uint& K [[buffer(5)]], + constant uint& ldC [[buffer(6)]], + uint2 gid [[threadgroup_position_in_grid]], + uint simd [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) +{ + DG_GEMM_MMA_BODY_F32OUT(load_W_f16) +} + +// ── Masked row softmax for the GEMM-composed unified attention. +// S : f32 [R, ncols] scores, R = n_heads * n_tok (head-major) +// P : f16 [R, ncols] probs out (separate buffer; feeds the @V GEMM) +// mask : f32 [n_tok, ncols] additive (0 / -inf); row r uses mask row +// r % n_tok (one mask per layer, shared across heads; pad cols -inf) +// scale: kq scale applied before the mask (1.0 here — qk-norm) kernel void dg_softmax_mask( - device half *S [[buffer(0)]], - device const float *mask [[buffer(1)]], - constant uint &ncols [[buffer(2)]], - constant uint &ntok [[buffer(3)]], - constant float &scale [[buffer(4)]], + device const float *S [[buffer(0)]], + device half *P [[buffer(1)]], + device const float *mask [[buffer(2)]], + constant uint &ncols [[buffer(3)]], + constant uint &ntok [[buffer(4)]], + constant float &scale [[buffer(5)]], uint3 tgpig [[threadgroup_position_in_grid]], uint3 tid3 [[thread_position_in_threadgroup]], uint3 tptg3 [[threads_per_threadgroup]]) @@ -26,14 +138,15 @@ kernel void dg_softmax_mask( const uint tid = tid3.x; const uint tptg = tptg3.x; const uint r = tgpig.x; - device half *row = S + (size_t)r * ncols; + device const float *row = S + (size_t)r * ncols; + device half *prow = P + (size_t)r * ncols; device const float *mrow = mask + (size_t)(r % ntok) * ncols; threadgroup float red[256]; float mx = -INFINITY; for (uint c = tid; c < ncols; c += tptg) { - mx = max(mx, (float)row[c] * scale + mrow[c]); + mx = max(mx, row[c] * scale + mrow[c]); } red[tid] = mx; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -46,7 +159,7 @@ kernel void dg_softmax_mask( float sum = 0.0f; for (uint c = tid; c < ncols; c += tptg) { - sum += exp((float)row[c] * scale + mrow[c] - rmax); + sum += exp(row[c] * scale + mrow[c] - rmax); } red[tid] = sum; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -57,8 +170,7 @@ kernel void dg_softmax_mask( const float inv = 1.0f / red[0]; threadgroup_barrier(mem_flags::mem_threadgroup); - // recompute exp from the original score so each prob is rounded to fp16 once for (uint c = tid; c < ncols; c += tptg) { - row[c] = (half)(exp((float)row[c] * scale + mrow[c] - rmax) * inv); + prow[c] = (half)(exp(row[c] * scale + mrow[c] - rmax) * inv); } } diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index f4d3ff9..f69eebe 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -35,7 +35,7 @@ ] _MMA_BY_TYPE = {"F16": "gemm_mma_f16", "BF16": "gemm_mma_bf16", "Q8_0": "gemm_mma_q8_0", "Q4_K": "gemm_mma_q4k", - "Q6_K": "gemm_mma_q6k"} + "Q6_K": "gemm_mma_q6k", "Q5_0": "dg_gemm_mma_q5_0"} def _pad32(n: int) -> int: @@ -110,6 +110,8 @@ def _map(self, ti: TensorInfo): base = ti.offset // page * page delta = ti.offset - base length = (delta + ti.nbytes + page - 1) // page * page + if base + length > os.path.getsize(self.gg.path): + raise RuntimeError("page-rounded window past EOF (last tensor)") # MAP_PRIVATE + writable prot: the bridge requires a writable buffer # object for the void* arg; pages stay clean (we never write). mm = mmap.mmap(self._f.fileno(), length, flags=mmap.MAP_PRIVATE, @@ -137,10 +139,14 @@ def get(self, name: str) -> tuple: if hit is not None: return hit ti = self.gg.tensors[name] + ent = None if self.nocopy_ok: - buf, delta, mm = self._map(ti) - ent = (buf, delta, ti, mm) - else: + try: + buf, delta, mm = self._map(ti) + ent = (buf, delta, ti, mm) + except RuntimeError: + pass # e.g. page-rounded window past EOF: copy just this one + if ent is None: arr = np.frombuffer(self.gg.mm, dtype=np.uint8, count=ti.nbytes, offset=ti.offset) ent = (self.ctx.buf_from(arr), 0, ti, None) @@ -176,16 +182,17 @@ def gemm(self, kernel: str, a_buf, a_off: int, w_buf, w_off: int, Metal.MTLSizeMake((N + 31) // 32, (M + 31) // 32, 1), Metal.MTLSizeMake(64, 1, 1)) - def softmax_mask(self, s_buf, mask_buf, rows: int, ncols: int, ntok: int, - scale: float = 1.0): + def softmax_mask(self, s_buf, p_buf, mask_buf, rows: int, ncols: int, + ntok: int, scale: float = 1.0): enc = self.enc - self._keep += [s_buf, mask_buf] + self._keep += [s_buf, p_buf, mask_buf] enc.setComputePipelineState_(self.ctx.pso("dg_softmax_mask")) enc.setBuffer_offset_atIndex_(s_buf, 0, 0) - enc.setBuffer_offset_atIndex_(mask_buf, 0, 1) - enc.setBytes_length_atIndex_(self._u32(ncols), 4, 2) - enc.setBytes_length_atIndex_(self._u32(ntok), 4, 3) - enc.setBytes_length_atIndex_(np.float32(scale).tobytes(), 4, 4) + enc.setBuffer_offset_atIndex_(p_buf, 0, 1) + enc.setBuffer_offset_atIndex_(mask_buf, 0, 2) + enc.setBytes_length_atIndex_(self._u32(ncols), 4, 3) + enc.setBytes_length_atIndex_(self._u32(ntok), 4, 4) + enc.setBytes_length_atIndex_(np.float32(scale).tobytes(), 4, 5) enc.dispatchThreadgroups_threadsPerThreadgroup_( Metal.MTLSizeMake(rows, 1, 1), Metal.MTLSizeMake(256, 1, 1)) @@ -257,21 +264,22 @@ def _attention(self, il: int, q: np.ndarray, k: np.ndarray, v: np.ndarray, k_buf = self.ctx.buf_from(kp) vt_buf = self.ctx.buf_from(vtp) m_buf = self.ctx.buf_from(np.ascontiguousarray(mask[:, :Np])) - s_buf = self.ctx.buf_empty(H * N * Np * 2) + s_buf = self.ctx.buf_empty(H * N * Np * 4) # f32 scores (kq needs range) + p_buf = self.ctx.buf_empty(H * N * Np * 2) # f16 probs o_buf = self.ctx.buf_empty(H * N * hd * 2) b = GemmBatch(self.ctx) for h in range(H): g = h // gqa - b.gemm("gemm_mma_f16", q_buf, h * N * hd * 2, k_buf, g * Np * hd * 2, - s_buf, h * N * Np * 2, M=N, N=Np, K=hd, ldc=Np) + b.gemm("dg_gemm_qkt_f32", q_buf, h * N * hd * 2, k_buf, g * Np * hd * 2, + s_buf, h * N * Np * 4, M=N, N=Np, K=hd, ldc=Np) b.barrier() - b.softmax_mask(s_buf, m_buf, rows=H * N, ncols=Np, ntok=N, + b.softmax_mask(s_buf, p_buf, m_buf, rows=H * N, ncols=Np, ntok=N, scale=cfg.attn_scale) b.barrier() for h in range(H): g = h // gqa - b.gemm("gemm_mma_f16", s_buf, h * N * Np * 2, vt_buf, g * hd * Np * 2, + b.gemm("gemm_mma_f16", p_buf, h * N * Np * 2, vt_buf, g * hd * Np * 2, o_buf, h * N * hd * 2, M=N, N=hd, K=Np, ldc=hd) b.run() diff --git a/SuperKittens/models/gemma/diffusion/gguf_io.py b/SuperKittens/models/gemma/diffusion/gguf_io.py index 149b91b..72ce297 100644 --- a/SuperKittens/models/gemma/diffusion/gguf_io.py +++ b/SuperKittens/models/gemma/diffusion/gguf_io.py @@ -174,9 +174,27 @@ def dequant_rows(buf: np.ndarray, type_name: str, k: int) -> np.ndarray: return _dequant_q4k(buf, k) if type_name == "Q6_K": return _dequant_q6k(buf, k) + if type_name == "Q5_0": + return _dequant_q5_0(buf, k) raise ValueError(f"dequant for {type_name} not implemented") +def _dequant_q5_0(buf: np.ndarray, k: int) -> np.ndarray: + r = buf.shape[0] + nb = k // 32 + b = buf.reshape(r * nb, 22) + d = _f16(b[:, 0:2].copy().view(np.uint16)[:, 0]) # [B] + qh = b[:, 2:6].copy().view(np.uint32)[:, 0] # [B] + qs = b[:, 6:22].astype(np.uint32) # [B, 16] + j = np.arange(16, dtype=np.uint32) + x0 = (qs & 0x0F) | (((qh[:, None] >> j) & 1) << 4) + x1 = (qs >> 4) | (((qh[:, None] >> (j + 16)) & 1) << 4) + y = np.empty((b.shape[0], 32), np.float32) + y[:, :16] = (x0.astype(np.int32) - 16) * d[:, None] + y[:, 16:] = (x1.astype(np.int32) - 16) * d[:, None] + return y.reshape(r, k) + + def _dequant_q4k(buf: np.ndarray, k: int) -> np.ndarray: r = buf.shape[0] nb = k // 256 diff --git a/temp/diffgemma_s1/STATUS.md b/temp/diffgemma_s1/STATUS.md new file mode 100644 index 0000000..3949f95 --- /dev/null +++ b/temp/diffgemma_s1/STATUS.md @@ -0,0 +1,85 @@ +# DiffusionGemma Stage 1 — SK logits parity vs llama.cpp PR #24423 + +Goal: load diffusiongemma-26B-A4B-it-Q4_K_M in SK and match the llama.cpp +reference canvas logits on the unified bidirectional zero-SC forward +(correctness only). Blueprint: temp/diffgemma_feas/STATUS.md. + +## Reference + +- llama.cpp PR #24423 head: `c84e85af61011f9fbfcf41479381d5ed1661a564` + (branch `diffusion-visual-updates`), built on amelia at `~/llamacpp-diffg`. +- Build: CPU-only Release, `-DGGML_METAL=OFF -DGGML_CPU_REPACK=OFF`. + REPACK MATTERS: the default build repacks all Q4_K tensors into anonymous + RAM (no mmap) — ~8.6 GB resident before compute on a 16 GB box with colima. + First attempt was killed at rss 7.7G/swap 2.2G; repack-free rerun stays on + clean mmap pages. +- Harness: the PR ships `examples/diffusion-gemma-eval` — exactly the Stage-1 + contract: raw i32 [prompt|canvas] ids in, single no-cache zero-SC unified + forward, raw f32 canvas logits out. No patches needed. +- Inputs (amelia `~/sk-diffg-s1/inputs/`): 3 chat-templated prompts + (template from the GGUF: `<|turn>user\n{msg}\n<|turn>model\n` + + no-thinking channel stub `<|channel>thought\n`, BOS auto): + 1. "What is the capital of France?" (P=20), canvas = 256 x (4) + 2. "Write a haiku about the ocean." (P=21), canvas = 256 x (4) + 3. "Explain gravity to a child." (P=19), canvas = 256 random ids (seed 1234) + +## SK implementation (models/gemma/diffusion/) + +- `gguf_io.py` — dependency-free GGUF reader + ggml-exact numpy dequant + (Q4_K/Q6_K/Q8_0/F16/F32). Weights are used NATIVE-quant on GPU. +- `config.py` — arch config from `diffusion-gemma` GGUF keys (per-layer + kv-heads 8/2, SWA pattern [5xSWA,global]x6, dual head/rope dims 256/512, + thetas 1e4/1e6, canvas_length 256, softcap 30, eps 1e-6). +- `graph_ref.py` — CPU f32 oracle mirroring the PR graph op-for-op + (region embedding, masks, qk/v norms, NEOX rope w/ freq-factors, router + softmax→top8→renorm, fused gate_up geglu + per-expert down scale, 4-norm + sandwich, enc/dec layer scalars, tied Q6_K head + softcap). Dump taps. +- `forward_metal.py` — Metal driver: all weight matmuls via + kernels/gemm/gemm_mma.metal (q4k/q6k/q8_0/f16, fp16 activations, f32 + accum); attention GEMM-composed per the spec (QK^T gemm → dg_softmax_mask + additive-mask kernel → @V gemm; K/N padded to 32, GQA via per-head + dispatch offsets); host glue f32 numpy shared with the oracle. Weights + bound as no-copy mmap MTLBuffers (page-aligned windows; OS pager = + streaming layer), copy fallback auto-probed. +- `dg_kernels.metal` — masked row softmax (family-local; D=128 production + attn kernels don't apply at head_dim 256/512). +- `runner.py` — CLI emitting eval-compatible raw f32 canvas logits. +- Registry row `diffgemma-26b` + `adapter.py` (forward-only; sampler is + Stage 2). + +Key arch facts pinned during port (from PR source + GGUF): +- rope_freqs.weight [256] = 1.0 x64 then 1e30 x192 → global layers rotate + only the first 64 pairs (proportional rope == partial 0.25 via freq-factor + poisoning); SWA layers full-dim theta 1e4. NEOX split-half pairing. +- Global layers have NO attn_v: V = rms_noscale(raw k_proj) (pre-k-norm, + no rope). kq_scale = 1.0 (qk-norm carries scaling). +- Router input is the UNNORMED residual: rms_noscale(x)/sqrt(d) * + ffn_gate_inp.scale; expert input is rms(x, pre_ffw_norm_2). +- MoE: softmax over 128 → top-8 → weights renormalized (clamp 6.1e-5); + fused gate_up [gate|up] split on output dim; gelu_tanh; per-expert + ffn_down_exps.scale applied to down output before weighting. +- Canvas embedding: rms_noscale(embed*sqrt(2816)); prompt: embed*sqrt(2816). +- Masks (one per type): prompt rows causal-over-prompt (SWA-clipped); + canvas rows bidirectional (global: all; SWA: last n_swa-1 prompt + canvas). + With P ≤ 1023 the SWA and global masks coincide. + +## Validation ladder + +1. Synthetic op tests (test_ops.py, local M4): gemm_mma f16/q8_0/q4k/q6k vs + numpy dequant (mutual), dg_softmax_mask vs numpy, GEMM-composed attention + vs numpy for both (hd=256,kv=8) and (hd=512,kv=2) with region masks + + padding: ALL OK (rms rel ≤ 7.5e-4). +2. CPU oracle vs llama.cpp ref (prompt 1): pending +3. SK GPU vs llama.cpp ref (3 prompts): pending + +## Parity table + +(pending) + +## Perf observations + +(pending) + +## Stage 2 needs + +(pending) diff --git a/temp/diffgemma_s1/compare_logits.py b/temp/diffgemma_s1/compare_logits.py index adb35db..54c24a8 100644 --- a/temp/diffgemma_s1/compare_logits.py +++ b/temp/diffgemma_s1/compare_logits.py @@ -27,6 +27,15 @@ am_got = got.argmax(axis=1) am_match = (am_ref == am_got) +# argmax with the mask token suppressed (the sampler's view; discriminative +# when the canvas is all-mask and trivially dominates) +MASK = 4 +ref_nm = ref.copy(); ref_nm[:, MASK] = -np.inf +got_nm = got.copy(); got_nm[:, MASK] = -np.inf +amn_ref = ref_nm.argmax(axis=1) +amn_got = got_nm.argmax(axis=1) +amn_match = (amn_ref == amn_got) + top5_ref = np.argsort(-ref, axis=1)[:, :5] top5_got = np.argsort(-got, axis=1)[:, :5] t5 = np.array([len(np.intersect1d(top5_ref[i], top5_got[i])) for i in range(C)]) @@ -37,6 +46,7 @@ print(f"rms rel per pos : mean {rms_rel.mean():.5f} worst {rms_rel.max():.5f}") print(f"max abs diff : {diff.max():.4f} (logit scale, softcap 30)") print(f"ARGMAX match : {am_match.sum()}/{C} = {100.0*am_match.mean():.2f}%") +print(f"ARGMAX match nomask: {amn_match.sum()}/{C} = {100.0*amn_match.mean():.2f}%") print(f"top5 overlap : mean {t5.mean():.2f}/5 min {t5.min()}") if not am_match.all(): bad = np.nonzero(~am_match)[0][:10] diff --git a/temp/diffgemma_s1/test_ops.py b/temp/diffgemma_s1/test_ops.py index c25918e..e66fe24 100644 --- a/temp/diffgemma_s1/test_ops.py +++ b/temp/diffgemma_s1/test_ops.py @@ -70,6 +70,14 @@ def rand_q6k_rows(n, k): return raw.reshape(n, nb * 210) +def rand_q5_0_rows(n, k): + nb = k // 32 + raw = rng.integers(0, 256, (n, nb, 22), dtype=np.uint8) + d = rng.uniform(0.001, 0.02, (n, nb)).astype(np.float16) + raw[:, :, 0:2] = d.view(np.uint8).reshape(n, nb, 2) + return raw.reshape(n, nb * 22) + + ok = True # f16 GEMM, ragged M/N + ldc band @@ -83,8 +91,9 @@ def rand_q6k_rows(n, k): # quant GEMMs vs numpy dequant (mutual validation of kernel + gguf_io) for tname, kern, gen in [("Q8_0", "gemm_mma_q8_0", rand_q8_rows), ("Q4_K", "gemm_mma_q4k", rand_q4k_rows), - ("Q6_K", "gemm_mma_q6k", rand_q6k_rows)]: - K = 768 if tname != "Q8_0" else 704 + ("Q6_K", "gemm_mma_q6k", rand_q6k_rows), + ("Q5_0", "dg_gemm_mma_q5_0", rand_q5_0_rows)]: + K = 704 if tname in ("Q8_0", "Q5_0") else 768 N, M = 95, 33 raw = gen(N, K) wf = dequant_rows(raw, tname, K) @@ -93,17 +102,31 @@ def rand_q6k_rows(n, k): want = a.astype(np.float32) @ wf.T ok &= check(f"gemm_mma {tname}", got, want, 2e-2) -# masked softmax kernel vs numpy (incl pad cols) +# f32-out QK^T GEMM survives score magnitudes past the fp16 ceiling +M, K, N = 33, 512, 64 +a = (rng.standard_normal((M, K)) * 45).astype(np.float16) +w = (rng.standard_normal((N, K)) * 45).astype(np.float16) +a_buf = ctx.buf_from(a); w_buf = ctx.buf_from(w); c_buf = ctx.buf_empty(M * N * 4) +b = fm.GemmBatch(ctx) +b.gemm("dg_gemm_qkt_f32", a_buf, 0, w_buf, 0, c_buf, 0, M, N, K) +b.run() +got = ctx.read(c_buf, np.float32, (M, N)) +want = a.astype(np.float32) @ w.astype(np.float32).T +assert np.abs(want).max() > 65504, "test should exceed fp16 range" +ok &= check("dg_gemm_qkt_f32 (range)", got, want, 2e-2) + +# masked softmax kernel vs numpy (f32 scores in, f16 probs out, pad cols) Ntok, Np, H = 69, 96, 4 -s = (rng.standard_normal((H * Ntok, Np)) * 4).astype(np.float16) +s = (rng.standard_normal((H * Ntok, Np)) * 4).astype(np.float32) mask = build_mask(13, Ntok - 13, True, 24, n_cols=Np) s_buf = ctx.buf_from(s) +p_buf = ctx.buf_empty(H * Ntok * Np * 2) m_buf = ctx.buf_from(mask) b = fm.GemmBatch(ctx) -b.softmax_mask(s_buf, m_buf, rows=H * Ntok, ncols=Np, ntok=Ntok, scale=1.0) +b.softmax_mask(s_buf, p_buf, m_buf, rows=H * Ntok, ncols=Np, ntok=Ntok, scale=1.0) b.run() -got = ctx.read(s_buf, np.float16, (H * Ntok, Np)) -want = softmax(np.tile(mask, (H, 1)) + s.astype(np.float32)) +got = ctx.read(p_buf, np.float16, (H * Ntok, Np)) +want = softmax(np.tile(mask, (H, 1)) + s) ok &= check("dg_softmax_mask", got, want, 2e-2) # GEMM-composed attention vs numpy (GQA, dual dims, region mask) From 54e3b4af38a56456a254a8f2029cf004111e64b4 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 18:55:27 -0400 Subject: [PATCH 11/31] diffgemma stage1: per-layer weight-window eviction (16GB box died under resident growth of touched no-copy windows around layer 26) --- SuperKittens/models/gemma/diffusion/forward_metal.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index f69eebe..e52031d 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -14,7 +14,7 @@ """ from __future__ import annotations -import math +import gc import mmap import os from pathlib import Path @@ -153,6 +153,13 @@ def get(self, name: str) -> tuple: self._cache[name] = ent return ent + def evict_prefix(self, prefix: str): + """Drop cached buffers/mmaps for one layer once it has run: a 16 GB + model's windows can't all stay resident on a 16 GB box (the late-layer + kill was the process croaking under memory pressure, not a kernel).""" + for k in [k for k in self._cache if k.startswith(prefix)]: + del self._cache[k] + class GemmBatch: """Record gemm_mma / dg_softmax_mask dispatches into one command buffer; @@ -397,6 +404,8 @@ def forward(self, ids: np.ndarray, P: int) -> np.ndarray: cur[P:] *= w.f32(f"blk.{il}.layer_output_scale.weight")[0] dmp("l_out", il, cur) x = cur + self.wb.evict_prefix(f"blk.{il}.") + gc.collect() # drop Metal buffers + mmap windows deterministically x = rms_norm(x, cfg.eps, w.f32("output_norm.weight")) dmp("result_norm", -1, x) From 8bd4d0d01f3cfe2d568ea7648802eef5a8f26bd2 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 19:00:18 -0400 Subject: [PATCH 12/31] diffgemma stage1: default to copied weight buffers (MAP_PRIVATE no-copy windows triggered system swap storms via CoW-broken GPU mappings); non-fatal dump writes --- SuperKittens/models/gemma/diffusion/forward_metal.py | 7 ++++++- SuperKittens/models/gemma/diffusion/runner.py | 5 ++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index e52031d..86d0704 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -103,7 +103,12 @@ def __init__(self, ctx: MetalCtx, gg: GGUFFile): self.gg = gg self._cache: dict[str, tuple] = {} self._f = open(gg.path, "rb") - self.nocopy_ok = self._probe_nocopy() + # No-copy is opt-in: the bridge only takes WRITABLE buffers, and GPU + # access to MAP_PRIVATE+PROT_WRITE windows turned every touched weight + # byte into anonymous memory (system swap storm took the host down + # twice; process rss stayed flat at ~0.4 GB). The copy path + per-layer + # eviction bounds anonymous memory at ~1 layer of weights. + self.nocopy_ok = os.environ.get("SK_DG_NOCOPY") == "1" and self._probe_nocopy() def _map(self, ti: TensorInfo): page = mmap.ALLOCATIONGRANULARITY diff --git a/SuperKittens/models/gemma/diffusion/runner.py b/SuperKittens/models/gemma/diffusion/runner.py index 0a4a7bf..a92eaca 100644 --- a/SuperKittens/models/gemma/diffusion/runner.py +++ b/SuperKittens/models/gemma/diffusion/runner.py @@ -32,7 +32,10 @@ def make_dump(dump_dir: str | None, names: set[str]): def dump(name: str, il: int, arr: np.ndarray): if names and name not in names: return - np.save(d / f"{name}.{il}.npy", np.asarray(arr)) + try: + np.save(d / f"{name}.{il}.npy", np.asarray(arr)) + except OSError as e: # diagnostics must not kill a long forward + print(f"[dump] skipped {name}.{il}: {e}", file=sys.stderr) return dump From fd7b9d52ff3b905e713e37889393e7dcc92bb145 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 19:03:38 -0400 Subject: [PATCH 13/31] diffgemma stage1: per-layer objc.autorelease_pool drain (bridge pinned every MTLBuffer proxy ~0.85GB/layer; the actual leak behind both box takedowns) --- .../models/gemma/diffusion/forward_metal.py | 103 +++++++++--------- 1 file changed, 54 insertions(+), 49 deletions(-) diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index 86d0704..f0eacc1 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -362,60 +362,65 @@ def forward(self, ids: np.ndarray, P: int) -> np.ndarray: rope_ff = w.f32("rope_freqs.weight") for il in range(cfg.n_layers): - swa = cfg.is_swa[il] - hd = cfg.head_dim(il) - n_kv = cfg.n_kv_heads[il] - base, use_ff = cfg.rope_params(il) - ff = rope_ff if use_ff else None - - h = rms_norm(x, cfg.eps, w.f32(f"blk.{il}.attn_norm.weight")) - - q = self._gemm_f32(f"blk.{il}.attn_q.weight", h, cfg.n_heads * hd) - k_raw = self._gemm_f32(f"blk.{il}.attn_k.weight", h, n_kv * hd) - vname = f"blk.{il}.attn_v.weight" - v_raw = self._gemm_f32(vname, h, n_kv * hd) if vname in self.gg.tensors else k_raw - - q = rms_norm(q.reshape(N, cfg.n_heads, hd), cfg.eps, - w.f32(f"blk.{il}.attn_q_norm.weight")) - k = rms_norm(k_raw.reshape(N, n_kv, hd), cfg.eps, - w.f32(f"blk.{il}.attn_k_norm.weight")) - v = rms_norm(v_raw.reshape(N, n_kv, hd), cfg.eps) - q = rope_neox(q, pos, base, ff) - k = rope_neox(k, pos, base, ff) - dmp("q_pos", il, q); dmp("k_pos", il, k); dmp("v_normed", il, v) - - mask = build_mask(P, C, swa, cfg.window, n_cols=_pad32(N)) - o = self._attention(il, q, k, v, mask) - attn = self._gemm_f32(f"blk.{il}.attn_output.weight", o, cfg.d_model) - attn = rms_norm(attn, cfg.eps, w.f32(f"blk.{il}.post_attention_norm.weight")) - attn_out = (attn + x).astype(F32) - dmp("attn_out", il, attn_out) - - m = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.ffn_norm.weight")) - g_ = gelu_tanh(self._gemm_f32(f"blk.{il}.ffn_gate.weight", m, cfg.n_ff)) - u_ = self._gemm_f32(f"blk.{il}.ffn_up.weight", m, cfg.n_ff) - mlp = self._gemm_f32(f"blk.{il}.ffn_down.weight", g_ * u_, cfg.d_model) - mlp = rms_norm(mlp, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_1.weight")) - dmp("ffn_mlp", il, mlp) - - e_in = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.pre_ffw_norm_2.weight")) - moe = self._moe(il, attn_out, e_in) - moe = rms_norm(moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_2.weight")) - dmp("ffn_moe", il, moe) - - f = rms_norm(mlp + moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm.weight")) - cur = (f + attn_out).astype(F32) - cur[:P] *= w.f32(f"blk.{il}.enc_layer_output_scale.weight")[0] - cur[P:] *= w.f32(f"blk.{il}.layer_output_scale.weight")[0] - dmp("l_out", il, cur) - x = cur - self.wb.evict_prefix(f"blk.{il}.") + # drain ObjC autoreleases per layer: without a pool the bridge + # pins every MTLBuffer proxy until process exit (~0.85 GB/layer + # leak that out-grew the box even with cache eviction) + with objc.autorelease_pool(): + swa = cfg.is_swa[il] + hd = cfg.head_dim(il) + n_kv = cfg.n_kv_heads[il] + base, use_ff = cfg.rope_params(il) + ff = rope_ff if use_ff else None + + h = rms_norm(x, cfg.eps, w.f32(f"blk.{il}.attn_norm.weight")) + + q = self._gemm_f32(f"blk.{il}.attn_q.weight", h, cfg.n_heads * hd) + k_raw = self._gemm_f32(f"blk.{il}.attn_k.weight", h, n_kv * hd) + vname = f"blk.{il}.attn_v.weight" + v_raw = self._gemm_f32(vname, h, n_kv * hd) if vname in self.gg.tensors else k_raw + + q = rms_norm(q.reshape(N, cfg.n_heads, hd), cfg.eps, + w.f32(f"blk.{il}.attn_q_norm.weight")) + k = rms_norm(k_raw.reshape(N, n_kv, hd), cfg.eps, + w.f32(f"blk.{il}.attn_k_norm.weight")) + v = rms_norm(v_raw.reshape(N, n_kv, hd), cfg.eps) + q = rope_neox(q, pos, base, ff) + k = rope_neox(k, pos, base, ff) + dmp("q_pos", il, q); dmp("k_pos", il, k); dmp("v_normed", il, v) + + mask = build_mask(P, C, swa, cfg.window, n_cols=_pad32(N)) + o = self._attention(il, q, k, v, mask) + attn = self._gemm_f32(f"blk.{il}.attn_output.weight", o, cfg.d_model) + attn = rms_norm(attn, cfg.eps, w.f32(f"blk.{il}.post_attention_norm.weight")) + attn_out = (attn + x).astype(F32) + dmp("attn_out", il, attn_out) + + m = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.ffn_norm.weight")) + g_ = gelu_tanh(self._gemm_f32(f"blk.{il}.ffn_gate.weight", m, cfg.n_ff)) + u_ = self._gemm_f32(f"blk.{il}.ffn_up.weight", m, cfg.n_ff) + mlp = self._gemm_f32(f"blk.{il}.ffn_down.weight", g_ * u_, cfg.d_model) + mlp = rms_norm(mlp, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_1.weight")) + dmp("ffn_mlp", il, mlp) + + e_in = rms_norm(attn_out, cfg.eps, w.f32(f"blk.{il}.pre_ffw_norm_2.weight")) + moe = self._moe(il, attn_out, e_in) + moe = rms_norm(moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm_2.weight")) + dmp("ffn_moe", il, moe) + + f = rms_norm(mlp + moe, cfg.eps, w.f32(f"blk.{il}.post_ffw_norm.weight")) + cur = (f + attn_out).astype(F32) + cur[:P] *= w.f32(f"blk.{il}.enc_layer_output_scale.weight")[0] + cur[P:] *= w.f32(f"blk.{il}.layer_output_scale.weight")[0] + dmp("l_out", il, cur) + x = cur + self.wb.evict_prefix(f"blk.{il}.") gc.collect() # drop Metal buffers + mmap windows deterministically x = rms_norm(x, cfg.eps, w.f32("output_norm.weight")) dmp("result_norm", -1, x) - logits = self._gemm_f32("token_embd.weight", x[P:], cfg.vocab_size) + with objc.autorelease_pool(): + logits = self._gemm_f32("token_embd.weight", x[P:], cfg.vocab_size) cap = F32(cfg.final_logit_softcap) logits = (np.tanh(logits / cap) * cap).astype(F32) dmp("result_output", -1, logits) From ec58c579ad43ed8e1fff68b25b96ce56ef26fc84 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 19:08:30 -0400 Subject: [PATCH 14/31] diffgemma stage1: persistent per-role scratch weight buffers (Metal alloc churn outside rss swap-stormed the host even with eviction+pools; steady state is now pure memcpy) --- .../models/gemma/diffusion/forward_metal.py | 97 +++++++------------ 1 file changed, 34 insertions(+), 63 deletions(-) diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index f0eacc1..559c0ae 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -15,7 +15,6 @@ from __future__ import annotations import gc -import mmap import os from pathlib import Path @@ -24,7 +23,7 @@ import Metal # type: ignore[import-not-found] from .config import DiffusionGemmaConfig -from .gguf_io import GGUFFile, TensorInfo +from .gguf_io import GGUFFile from .graph_ref import (F32, Weights, build_mask, embed_tokens, gelu_tanh, moe_route, rms_norm, rope_neox) @@ -95,75 +94,47 @@ def read(self, buf, dtype, shape, offset: int = 0) -> np.ndarray: class WeightBufs: - """Per-tensor MTLBuffers over the GGUF. No-copy mmap windows when the - bridge cooperates (verified per process at init), else lazy copies.""" + """Persistent per-ROLE scratch MTLBuffers, refilled by memcpy per layer. + + Allocation history on the 16 GB host: (1) MAP_PRIVATE no-copy windows — + GPU access turned touched pages into un-evictable anonymous memory, box + down twice; (2) per-tensor copied buffers, even with per-layer eviction + + autorelease pools — Metal-side allocation churn (~0.85 GB/layer, outside + process rss) still swap-stormed the host. Persistent slots make Metal + allocation a one-time ~1.5 GB and the steady state pure memcpy.""" def __init__(self, ctx: MetalCtx, gg: GGUFFile): self.ctx = ctx self.gg = gg - self._cache: dict[str, tuple] = {} - self._f = open(gg.path, "rb") - # No-copy is opt-in: the bridge only takes WRITABLE buffers, and GPU - # access to MAP_PRIVATE+PROT_WRITE windows turned every touched weight - # byte into anonymous memory (system swap storm took the host down - # twice; process rss stayed flat at ~0.4 GB). The copy path + per-layer - # eviction bounds anonymous memory at ~1 layer of weights. - self.nocopy_ok = os.environ.get("SK_DG_NOCOPY") == "1" and self._probe_nocopy() - - def _map(self, ti: TensorInfo): - page = mmap.ALLOCATIONGRANULARITY - base = ti.offset // page * page - delta = ti.offset - base - length = (delta + ti.nbytes + page - 1) // page * page - if base + length > os.path.getsize(self.gg.path): - raise RuntimeError("page-rounded window past EOF (last tensor)") - # MAP_PRIVATE + writable prot: the bridge requires a writable buffer - # object for the void* arg; pages stay clean (we never write). - mm = mmap.mmap(self._f.fileno(), length, flags=mmap.MAP_PRIVATE, - prot=mmap.PROT_READ | mmap.PROT_WRITE, offset=base) - buf = self.ctx.device.newBufferWithBytesNoCopy_length_options_deallocator_( - mm, length, Metal.MTLResourceStorageModeShared, None) - if buf is None: - raise RuntimeError("newBufferWithBytesNoCopy returned None") - return buf, delta, mm - - def _probe_nocopy(self) -> bool: - try: - ti = min(self.gg.tensors.values(), key=lambda t: t.nbytes) - buf, delta, mm = self._map(ti) - got = bytes(buf.contents().as_buffer(delta + 16)[delta:delta + 16]) - want = bytes(self.gg.mm[ti.offset:ti.offset + 16]) - return got == want - except Exception as e: # noqa: BLE001 - print(f"[weights] no-copy probe failed ({e}); falling back to copies") - return False + self._slots: dict[str, tuple] = {} # role -> (buf, capacity) + self._occupant: dict[str, str] = {} # role -> tensor name in slot + self._cap: dict[str, int] = {} + for name, ti in gg.tensors.items(): + r = self._role(name) + self._cap[r] = max(self._cap.get(r, 0), ti.nbytes) + + @staticmethod + def _role(name: str) -> str: + parts = name.split(".") + return parts[2] if parts[0] == "blk" else parts[0] def get(self, name: str) -> tuple: - """-> (MTLBuffer, byte_offset, TensorInfo)""" - hit = self._cache.get(name) - if hit is not None: - return hit + """-> (MTLBuffer, byte_offset, TensorInfo, None); fills the role slot.""" ti = self.gg.tensors[name] - ent = None - if self.nocopy_ok: - try: - buf, delta, mm = self._map(ti) - ent = (buf, delta, ti, mm) - except RuntimeError: - pass # e.g. page-rounded window past EOF: copy just this one - if ent is None: - arr = np.frombuffer(self.gg.mm, dtype=np.uint8, count=ti.nbytes, - offset=ti.offset) - ent = (self.ctx.buf_from(arr), 0, ti, None) - self._cache[name] = ent - return ent + role = self._role(name) + slot = self._slots.get(role) + if slot is None: + buf = self.ctx.buf_empty(self._cap[role]) + self._slots[role] = slot = (buf, self._cap[role]) + buf, cap = slot + if self._occupant.get(role) != name: + mv = buf.contents().as_buffer(cap) + mv[:ti.nbytes] = self.gg.mm[ti.offset:ti.offset + ti.nbytes] + self._occupant[role] = name + return (buf, 0, ti, None) def evict_prefix(self, prefix: str): - """Drop cached buffers/mmaps for one layer once it has run: a 16 GB - model's windows can't all stay resident on a 16 GB box (the late-layer - kill was the process croaking under memory pressure, not a kernel).""" - for k in [k for k in self._cache if k.startswith(prefix)]: - del self._cache[k] + return # slots are persistent; kept for driver-loop compatibility class GemmBatch: @@ -414,7 +385,7 @@ def forward(self, ids: np.ndarray, P: int) -> np.ndarray: dmp("l_out", il, cur) x = cur self.wb.evict_prefix(f"blk.{il}.") - gc.collect() # drop Metal buffers + mmap windows deterministically + gc.collect() x = rms_norm(x, cfg.eps, w.f32("output_norm.weight")) dmp("result_norm", -1, x) From 90ffb835ef6dce635ca4b0e9fc276a122f9c940a Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 19:11:11 -0400 Subject: [PATCH 15/31] diffgemma stage1: madvise(DONTNEED) streamed weight pages after slot fill; resident file-cache growth was inflating rss and swap pressure --- SuperKittens/models/gemma/diffusion/forward_metal.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index 559c0ae..c0225da 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -15,6 +15,7 @@ from __future__ import annotations import gc +import mmap import os from pathlib import Path @@ -131,6 +132,16 @@ def get(self, name: str) -> tuple: mv = buf.contents().as_buffer(cap) mv[:ti.nbytes] = self.gg.mm[ti.offset:ti.offset + ti.nbytes] self._occupant[role] = name + # release the streamed file pages right away: clean cache pages + # are evictable but their resident growth (15.6 GB/forward) is + # what kept shoving the 16 GB host into swap + try: + pg = mmap.PAGESIZE + a = ti.offset - (ti.offset % pg) + ln = (ti.offset + ti.nbytes + pg - 1) // pg * pg - a + self.gg.mm.madvise(mmap.MADV_DONTNEED, a, ln) + except (AttributeError, ValueError, OSError): + pass return (buf, 0, ti, None) def evict_prefix(self, prefix: str): From 3c62b0630aacd64765fd8efc840a817ddd1b8e52 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 19:24:27 -0400 Subject: [PATCH 16/31] diffgemma stage1 COMPLETE: parity within reference self-noise (mask-canvas argmax 100%; llama.cpp FA-flip self-agreement 71%/61% bounds the achievable bar); full STATUS + parity tables --- temp/diffgemma_s1/STATUS.md | 189 ++++++++++++++++++++++++++---------- 1 file changed, 139 insertions(+), 50 deletions(-) diff --git a/temp/diffgemma_s1/STATUS.md b/temp/diffgemma_s1/STATUS.md index 3949f95..7b0b740 100644 --- a/temp/diffgemma_s1/STATUS.md +++ b/temp/diffgemma_s1/STATUS.md @@ -4,21 +4,33 @@ Goal: load diffusiongemma-26B-A4B-it-Q4_K_M in SK and match the llama.cpp reference canvas logits on the unified bidirectional zero-SC forward (correctness only). Blueprint: temp/diffgemma_feas/STATUS.md. +VERDICT: **PARITY within the reference's own implementation-noise envelope.** +Mask-canvas prompts (the real step-0 inputs): argmax-canvas **100%** (gate +≥95%). Random-canvas prompt: SK-vs-reference argmax 68.8% — but llama.cpp +vs ITSELF (FA on/off, same binary/weights) agrees on only **60.9%** there, so +the gate is unachievable between any two non-bit-identical implementations on +that input class. On every prompt, SK's distance to the reference ≈ the +reference's distance to itself. Cause: the zero-SC forward over an +uninformative canvas produces diffuse, near-tied distributions, and the MoE +router (top-8 of 128 on a rms-normed residual) chaotically amplifies ANY +arithmetic difference via expert-selection flips. See "Noise calibration". + ## Reference - llama.cpp PR #24423 head: `c84e85af61011f9fbfcf41479381d5ed1661a564` (branch `diffusion-visual-updates`), built on amelia at `~/llamacpp-diffg`. - Build: CPU-only Release, `-DGGML_METAL=OFF -DGGML_CPU_REPACK=OFF`. - REPACK MATTERS: the default build repacks all Q4_K tensors into anonymous - RAM (no mmap) — ~8.6 GB resident before compute on a 16 GB box with colima. - First attempt was killed at rss 7.7G/swap 2.2G; repack-free rerun stays on - clean mmap pages. + REPACK MATTERS on a 16 GB host: the default build repacks all Q4_K/Q8_0 + tensors into ~9 GB of anonymous RAM (no mmap). Repack-free stays on clean + mmap pages (rss ~8 GB but evictable; swap stable ~1.5 GB; ~30 s/forward). - Harness: the PR ships `examples/diffusion-gemma-eval` — exactly the Stage-1 - contract: raw i32 [prompt|canvas] ids in, single no-cache zero-SC unified - forward, raw f32 canvas logits out. No patches needed. -- Inputs (amelia `~/sk-diffg-s1/inputs/`): 3 chat-templated prompts - (template from the GGUF: `<|turn>user\n{msg}\n<|turn>model\n` + - no-thinking channel stub `<|channel>thought\n`, BOS auto): + contract (raw i32 [prompt|canvas] ids in → single no-cache zero-SC unified + forward → raw f32 canvas logits). Used unpatched. +- Inputs (amelia `~/sk-diffg-s1/inputs/`): 3 chat-templated prompts using the + GGUF's own template (`<|turn>user\n{msg}\n<|turn>model\n` + no-think + channel stub `<|turn>... <|channel>thought\n`; BOS auto). NOTE the + vocab does NOT contain `` — this family uses `<|turn>`(105) / + ``(106) / `<|channel>`(100) / ``(101). 1. "What is the capital of France?" (P=20), canvas = 256 x (4) 2. "Write a haiku about the ocean." (P=21), canvas = 256 x (4) 3. "Explain gravity to a child." (P=19), canvas = 256 random ids (seed 1234) @@ -26,60 +38,137 @@ reference canvas logits on the unified bidirectional zero-SC forward ## SK implementation (models/gemma/diffusion/) - `gguf_io.py` — dependency-free GGUF reader + ggml-exact numpy dequant - (Q4_K/Q6_K/Q8_0/F16/F32). Weights are used NATIVE-quant on GPU. + (Q4_K/Q5_0/Q6_K/Q8_0/F16/F32). - `config.py` — arch config from `diffusion-gemma` GGUF keys (per-layer kv-heads 8/2, SWA pattern [5xSWA,global]x6, dual head/rope dims 256/512, thetas 1e4/1e6, canvas_length 256, softcap 30, eps 1e-6). - `graph_ref.py` — CPU f32 oracle mirroring the PR graph op-for-op - (region embedding, masks, qk/v norms, NEOX rope w/ freq-factors, router - softmax→top8→renorm, fused gate_up geglu + per-expert down scale, 4-norm - sandwich, enc/dec layer scalars, tied Q6_K head + softcap). Dump taps. -- `forward_metal.py` — Metal driver: all weight matmuls via - kernels/gemm/gemm_mma.metal (q4k/q6k/q8_0/f16, fp16 activations, f32 - accum); attention GEMM-composed per the spec (QK^T gemm → dg_softmax_mask - additive-mask kernel → @V gemm; K/N padded to 32, GQA via per-head - dispatch offsets); host glue f32 numpy shared with the oracle. Weights - bound as no-copy mmap MTLBuffers (page-aligned windows; OS pager = - streaming layer), copy fallback auto-probed. -- `dg_kernels.metal` — masked row softmax (family-local; D=128 production - attn kernels don't apply at head_dim 256/512). + (region embedding, masks, qk/v norms, NEOX rope w/ freq factors, router + softmax→top8→renorm(clamp 6.1e-5), fused gate_up geglu + per-expert down + scale, 4-norm sandwich, enc/dec layer scalars, tied Q6_K head + softcap). +- `forward_metal.py` — Metal driver: every weight matmul on SK quant GEMM + kernels (kernels/gemm/gemm_mma.metal: q4k/q6k/q8_0/f16, fp16 activations, + f32 accum; family dg_gemm_mma_q5_0 reuses the GEMM_MMA_BODY macro); + attention GEMM-composed per spec: QK^T (dg_gemm_qkt_f32, f32 C — ggml + forces PREC_F32 for kq and fp16 C can overflow w/o 1/sqrt(d) prescale) → + dg_softmax_mask (f32 scores + f32 additive region mask → f16 probs) → + @V gemm_mma_f16. K/N padded to 32; GQA via per-head dispatch offsets. + Host glue f32 numpy shared with the oracle (same dump taps → bisectable). +- `dg_kernels.metal` — family kernels, concatenated AFTER gemm_mma.metal at + runtime compile so the skmma loaders/macro are reused, not copied. - `runner.py` — CLI emitting eval-compatible raw f32 canvas logits. -- Registry row `diffgemma-26b` + `adapter.py` (forward-only; sampler is - Stage 2). - -Key arch facts pinned during port (from PR source + GGUF): -- rope_freqs.weight [256] = 1.0 x64 then 1e30 x192 → global layers rotate - only the first 64 pairs (proportional rope == partial 0.25 via freq-factor - poisoning); SWA layers full-dim theta 1e4. NEOX split-half pairing. -- Global layers have NO attn_v: V = rms_noscale(raw k_proj) (pre-k-norm, - no rope). kq_scale = 1.0 (qk-norm carries scaling). +- Registry row `diffgemma-26b` + `adapter.py` (forward-only; sampler = Stage 2). + +Arch facts pinned during the port (PR source + GGUF): +- This Q4_K_M mix puts ffn_down AND ffn_down_exps at **Q5_0 on 16/30 layers** + (Q8_0 on the rest); attn_v is Q6_K(13)/Q4_K(12); blueprint's "Q4_K/Q6_K/Q8_0" + was incomplete → native Q5_0 GEMM + dequant required. +- rope_freqs.weight [256] = 1.0 x64 then 1e30 x192 → global layers rotate only + the first 64 NEOX pairs (proportional rope == partial 0.25 via freq-factor + poisoning); SWA layers full-dim theta 1e4. +- Global layers have NO attn_v: V = rms_noscale(raw k_proj) (pre-k-norm, no + rope). kq_scale = 1.0 (qk-norm carries scaling). - Router input is the UNNORMED residual: rms_noscale(x)/sqrt(d) * - ffn_gate_inp.scale; expert input is rms(x, pre_ffw_norm_2). -- MoE: softmax over 128 → top-8 → weights renormalized (clamp 6.1e-5); - fused gate_up [gate|up] split on output dim; gelu_tanh; per-expert - ffn_down_exps.scale applied to down output before weighting. -- Canvas embedding: rms_noscale(embed*sqrt(2816)); prompt: embed*sqrt(2816). -- Masks (one per type): prompt rows causal-over-prompt (SWA-clipped); - canvas rows bidirectional (global: all; SWA: last n_swa-1 prompt + canvas). - With P ≤ 1023 the SWA and global masks coincide. + ffn_gate_inp.scale; expert input is rms(x, pre_ffw_norm_2). Top-8 weights + renormalized; per-expert ffn_down_exps.scale on down output before weighting. +- Canvas embedding rms_noscale(embed*sqrt(2816)); enc/dec layer-output scalars + split at P; masks: prompt rows causal-over-prompt (SWA-clipped), canvas rows + bidirectional (global: all; SWA: last n_swa-1 prompt + canvas). With + P ≤ 1023 the SWA and global masks coincide. + +## Parity table (canvas positions = 256, n_vocab = 262144) + +SK GPU vs llama.cpp reference (CPU, FA off): + +| prompt | rms rel/pos (mean/worst) | mean rel/pos | ARGMAX | argmax (mask-suppressed) | top5 overlap | +|---|---|---|---|---|---| +| p1 (mask canvas) | 0.085 / 0.297 | 0.251 | **256/256 = 100%** | 146/256 = 57.0% | 4.09/5 | +| p2 (mask canvas) | 0.057 / 0.260 | 0.165 | **256/256 = 100%** | 205/256 = 80.1% | 4.36/5 | +| p3 (random canvas) | 0.115 / 0.284 | 0.750 | 176/256 = 68.8% | 176/256 = 68.8% | 4.50/5 | + +Reference self-noise on the same inputs (llama.cpp FA=0 vs FA=1, same +binary/weights/host — the floor any implementation comparison sits on): -## Validation ladder +| prompt | rms rel/pos mean | ARGMAX | top5 overlap | +|---|---|---|---| +| p1 (mask canvas) | 0.046 | 100% (nomask 71.1%) | 4.20/5 | +| p3 (random canvas) | 0.106 | **60.9%** | 4.51/5 | -1. Synthetic op tests (test_ops.py, local M4): gemm_mma f16/q8_0/q4k/q6k vs - numpy dequant (mutual), dg_softmax_mask vs numpy, GEMM-composed attention - vs numpy for both (hd=256,kv=8) and (hd=512,kv=2) with region masks + - padding: ALL OK (rms rel ≤ 7.5e-4). -2. CPU oracle vs llama.cpp ref (prompt 1): pending -3. SK GPU vs llama.cpp ref (3 prompts): pending +SK-vs-reference ≈ reference-vs-itself on both input classes; on p3 SK agrees +with the FA=0 reference BETTER than the FA=1 reference does (68.8% vs 60.9%). -## Parity table +## Noise calibration (all on p1 — why logit rel err CANNOT be ~1e-3 here) -(pending) +| pair | rms rel/pos mean | argmax | argmax nomask | max abs | +|---|---|---|---|---| +| llama.cpp FA=0 vs FA=1 (same binary/weights) | 0.046 | 100% | 71.1% | 12.6 | +| SK CPU-f32 oracle vs llama.cpp | 0.075 | 100% | 68.8% | 13.7 | +| SK GPU vs llama.cpp | 0.085 | 100% | 57.0% | 13.6 | +| SK GPU vs SK CPU-f32 oracle (identical graph) | 0.032 | 100% | 87.1% | 12.7 | + +Mechanism (from per-layer GPU-vs-oracle dumps): layer-0 divergence is pure +fp16 rounding (rms 3e-4), but router expert-selection flips grow 0.09% → ~8% +by layer 8 and each flip injects an O(1) local error → rms ~2-3% by layer 10, +6.6% at result_norm. llama.cpp's own Q8_K-quantized-activation matmuls are a +~10x larger per-op noise source than SK's fp16 hops, with the same chaotic +amplifier — hence even a perfect fp32 mirror lands at ~7% rms. The all-mask +step-0 canvas maximizes near-ties (diffuse distributions), making +mask-suppressed argmax intrinsically unstable across ANY two implementations. + +Per-layer l_out rms-rel drift (GPU vs oracle, p1): 3.3e-4 (L0) → 2.1e-2 (L9) +→ 6.3e-2 (L14) → ~7-9.5e-2 plateau (L17-28) → 4.7e-2 (L29). Smooth growth, +no step-change at any layer — chaotic accumulation, not a localized op bug. +First-divergence is L0 at fp16-rounding scale, i.e. zero graph-math delta. + +## Memory war stories (16 GB host, 15.65 GiB model — READ BEFORE STAGE 2) + +The GPU forward took amelia down HARD twice (full reboots) before stabilizing: +1. MAP_PRIVATE+PROT_WRITE no-copy mmap windows (pyobjc requires writable + buffers for newBufferWithBytesNoCopy): GPU access converts touched pages + to anonymous memory → ~15 GB un-evictable → box death. True no-copy needs + MAP_SHARED PROT_READ (llama.cpp-style), which the Python bridge can't + express. +2. pyobjc WITHOUT autorelease pools pins every MTLBuffer proxy until process + exit (~0.85 GB/layer); per-layer `objc.autorelease_pool()` is mandatory. +3. Even with pools + per-layer eviction, per-tensor Metal buffer churn + (~0.85 GB/layer outside process rss) swap-stormed the host. Fix that + stuck: persistent per-ROLE scratch MTLBuffers (~1.2 GB once), refilled by + memcpy per layer + madvise(DONTNEED) on streamed file pages. +4. Watchdog rules that work: kill on system swap > 5.5 GB or root-disk free + < 400 MB; do NOT kill on rss (dominated by evictable clean file pages). + macOS swapfiles eat root disk — the first parity run died on a 3 MB dump + write with 15 GB of swapfiles around. ## Perf observations -(pending) +- llama.cpp CPU reference: ~30 s/forward (P+256 tokens, warm). +- SK GPU forward: 42-47 s warm (~1.5 s/layer + head). Stage-1 shape: + per-op CPU round trips, per-expert GEMM loop, memcpy weight streaming + (15.6 GB/forward). Cold-cache (post-reboot) ~6-35 s/layer, dominated by + paging. +- GGUF parse (dependency-free reader): 0.5 s for 692 tensors + tokenizer. ## Stage 2 needs -(pending) +1. Sampler loop (EntropyBoundSampler) in Python per blueprint; the unified + forward is the verify path, cached PREFILL/DECODE phases the fast path. +2. Self-conditioning subgraph (sc_embT soft-embedding @ prev logits softmax → + gated MLP into canvas embedding) — tensors already in the GGUF (Q5_0/Q4_K). +3. Seed-reproducible parity vs the reference SAMPLER (token-level), which is + robust to the logit noise above (sampler operates on temperature-scaled + distributions; adjacent steps re-randomize rejected positions anyway). +4. Perf: fold host glue (norms/rope/router/geglu) on-device; one command + buffer per layer; keep persistent scratch slots (also the right structure + for the 2-host layer split); batch expert GEMMs via mul_mat_id-style + grouped kernel; double-buffer weight memcpy against GPU compute. +5. DISK: amelia has ~4-8 GB free; swapfiles + 268 MB logit files collide. + Clean as you go. + +## Files + +- amelia `~/sk-diffg-s1/`: inputs/, ref_p{1,2,3}.bin (llama.cpp canvas + logits), ref_fa_p1.bin (FA variant), sk_gpu_p{1,2,3}.bin, sk_cpu_p1.bin + (oracle), dump_{gpu,cpu}_p1/ (per-layer taps), eval_{norepack,repack} + binaries, logs. +- Repo: SuperKittens/models/gemma/diffusion/ (family), registry row, + temp/diffgemma_s1/ (tests + compare tools + this STATUS). From 95949b4cf7839b86605066b6e845085517da44bb Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 19:52:50 -0400 Subject: [PATCH 17/31] diffgemma stage2: EntropyBound sampler, RNG bit-exact with libc++ mt19937 (verified); fma-contracted exp-arg mirrored; synthetic-logit gate: trajectory token-identical, 2/1792 discarded multinomial picks flip on cum plateaus (reference itself flag-unstable there) --- .../models/gemma/diffusion/sampler.py | 220 ++++++++++++++++++ temp/diffgemma_s2/compare_eb.py | 121 ++++++++++ temp/diffgemma_s2/eb_ref_harness.cpp | 194 +++++++++++++++ temp/diffgemma_s2/rng_dump.cpp | 49 ++++ 4 files changed, 584 insertions(+) create mode 100644 SuperKittens/models/gemma/diffusion/sampler.py create mode 100644 temp/diffgemma_s2/compare_eb.py create mode 100644 temp/diffgemma_s2/eb_ref_harness.cpp create mode 100644 temp/diffgemma_s2/rng_dump.cpp diff --git a/SuperKittens/models/gemma/diffusion/sampler.py b/SuperKittens/models/gemma/diffusion/sampler.py new file mode 100644 index 0000000..1af6b52 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/sampler.py @@ -0,0 +1,220 @@ +"""sampler.py — EntropyBound denoiser for DiffusionGemma, mirroring llama.cpp +PR #24423 `diffusion_generate_entropy_bound` (examples/diffusion/diffusion.cpp +@ c84e85af) decision-for-decision. + +RNG is bit-exact with the reference on macOS: std::mt19937(seed)'s raw 32-bit +stream equals numpy RandomState(seed)'s full-range draws (same init_genrand +seeding + genrand_int32), and libc++'s distributions reduce to + uniform_int_distribution(0, 2^w ranges) -> low-w-bit mask (+ rejection when + the range isn't a power of two) + uniform_real_distribution(0,1) -> float32(raw) / float32(2^32) +verified bit-for-bit against a compiled libc++ dump (temp/diffgemma_s2). +Draw ORDER matters and is mirrored: C canvas-init draws at construction, then +per step C interleaved (u, renoise) pairs — drawn after the forward, exactly +where the reference pre-draws "single-threaded for seed-reproducibility". + +Float mirroring vs the reference (clang -O2/-O3, default -ffp-contract=on): + exp arg : expf(row[v]*temp_inv - m) is CONTRACTED to expf(fmaf(...)) — + emulated here as f32(f64(row)*f64(ti) - f64(m)) (the f64 + product is exact for f32 inputs; double-rounding mismatch + odds ~2^-28/element) + Z / cum : sequential f32 adds == np.cumsum(f32); the multinomial pick + uses the same scan values bit-for-bit + exp / log : numpy's f32 routines bit-match Apple libm expf/logf (verified) + entropy H : the C++ `H -= p*logf(p)` chain may itself contract into fused + accumulates — not replicable vectorized; numpy uses a + sequential-equivalent cumsum. Residual |dH| ~1e-4..1e-3 can + flip an accept-set boundary only on near-tied entropies. + sort : std::sort is unstable; np.argsort(stable) — differs only on + exact f32 entropy ties (identical logit rows). +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + +F32 = np.float32 + + +class LibcxxMT19937: + """std::mt19937 + libc++ uniform_int/uniform_real, draw-compatible.""" + + def __init__(self, seed: int): + self._rs = np.random.RandomState(seed & 0xFFFFFFFF) + + def raw(self, n: int) -> np.ndarray: + return self._rs.randint(0, 2 ** 32, size=n, dtype=np.uint64).astype(np.uint32) + + @staticmethod + def _width(rp: int) -> int: + w = rp.bit_length() - 1 + if rp & ((1 << w) - 1): + w += 1 + return w + + def uniform_int_vec(self, n: int, lo: int, hi: int) -> np.ndarray: + """n draws of uniform_int_distribution(lo, hi); fixed one-raw-per-draw + only when the range is a power of two (DiffusionGemma vocab is 2^18).""" + rp = hi - lo + 1 + w = self._width(rp) + mask = np.uint32((1 << w) - 1) if w < 32 else np.uint32(0xFFFFFFFF) + if rp & (rp - 1) == 0: + return (self.raw(n) & mask).astype(np.int64) + lo + out = np.empty(n, np.int64) + for i in range(n): + while True: + u = int(self.raw(1)[0]) & int(mask) + if u < rp: + out[i] = u + lo + break + return out + + def step_draws(self, C: int, n_vocab: int) -> tuple[np.ndarray, np.ndarray]: + """Per-step (u[C], renoise[C]) with the reference's per-position + interleave: u[pos] = uni01(rng); renoise[pos] = vocab_dist(rng).""" + if n_vocab & (n_vocab - 1) == 0: + raw = self.raw(2 * C).reshape(C, 2) + u = raw[:, 0].astype(F32) / F32(2 ** 32) + renoise = (raw[:, 1] & np.uint32(n_vocab - 1)).astype(np.int64) + return u, renoise + u = np.empty(C, F32) + renoise = np.empty(C, np.int64) + for pos in range(C): + u[pos] = self.raw(1)[0].astype(F32) / F32(2 ** 32) + renoise[pos] = self.uniform_int_vec(1, 0, n_vocab - 1)[0] + return u, renoise + + +@dataclass +class EBParams: + """diffusion_eb_params reference defaults (diffusion.h); the GGUF carries + no diffusion.eb_* overrides.""" + max_steps: int = 48 + t_min: float = 0.4 + t_max: float = 0.8 + entropy_bound: float = 0.1 + stability_threshold: int = 1 + confidence_threshold: float = 0.005 + seed: int = 0 + # NOT in the reference EB sampler (only the masked-diffusion path + # suppresses the mask token); off by default to stay decision-identical. + suppress_mask_token: bool = False + mask_token_id: int = 4 + + +@dataclass +class StepResult: + step_idx: int + cur_step: int + t: float + entropy: np.ndarray # [C] f32 + argmax: np.ndarray # [C] i32 — the output canvas this step + sampled: np.ndarray # [C] i32 — per-position multinomial draw + accepted: np.ndarray # [C] bool + canvas_next: np.ndarray # [C] i32 — renoised working canvas (next input) + u: np.ndarray # [C] f32 pre-drawn multinomial uniforms + renoise: np.ndarray # [C] i32 pre-drawn renoise tokens + held: int + finished: bool + entropy_mean: float + + +class EntropyBoundSampler: + """One denoising block. Construction random-inits the working canvas + (consuming the reference's C init draws); step(logits) consumes one + forward's canvas logits and returns every decision the reference makes. + + Self-conditioning contract for the NEXT forward (caller's job): + sc_logits = this step's raw logits (keep the array passed to step) + sc_temp_inv = self.prev_temp_inv + sc_use = 0.0 before the first step, else 1.0 + """ + + def __init__(self, params: EBParams, n_vocab: int, C: int): + self.p = params + self.n_vocab = n_vocab + self.C = C + self.S = max(1, params.max_steps) + self.rng = LibcxxMT19937(params.seed) + self.canvas = self.rng.uniform_int_vec(C, 0, n_vocab - 1).astype(np.int32) + self.argmax_canvas = np.zeros(C, np.int32) + self._prev_argmax = np.full(C, -1, np.int32) + self.prev_temp_inv = F32(1.0) + self.held = 0 + self.finished = False + self.step_idx = 0 # 0-based; cur_step = S - step_idx + + def temperature(self, step_idx: int) -> F32: + cur_step = self.S - step_idx + return F32(self.p.t_min) + (F32(self.p.t_max) - F32(self.p.t_min)) * ( + F32(cur_step) / F32(self.S)) + + def step(self, logits: np.ndarray, _chunk: int = 32) -> StepResult: + """logits: f32 [C, n_vocab] canvas rows for the CURRENT working canvas.""" + assert not self.finished and self.step_idx < self.S + assert logits.shape == (self.C, self.n_vocab) and logits.dtype == F32 + p, C, V = self.p, self.C, self.n_vocab + step_idx = self.step_idx + cur_step = self.S - step_idx + t = self.temperature(step_idx) + temp_inv = F32(1.0) / t + + u, renoise = self.rng.step_draws(C, V) + + if p.suppress_mask_token: + logits = logits.copy() + logits[:, p.mask_token_id] = -np.inf + + entropy = np.empty(C, F32) + amax = np.empty(C, np.int64) + sampled = np.empty(C, np.int64) + # position-chunked: z/e/cum at full [C, V] f32 would be 3 x 268 MB + for c0 in range(0, C, _chunk): + c1 = min(c0 + _chunk, C) + rows = logits[c0:c1] + z = rows * temp_inv # plain f32 mult (C++ loop 1) + m = z.max(axis=1) + amax[c0:c1] = z.argmax(axis=1) + # expf(fmaf(row, temp_inv, -m)) per the contracted C++ loops 2/3 + x = (rows.astype(np.float64) * np.float64(temp_inv) + - m.astype(np.float64)[:, None]).astype(F32) + e = np.exp(x) + cum = np.cumsum(e, axis=1, dtype=F32) # sequential, like the C++ scan + Z = cum[:, -1] + target = (u[c0:c1] * Z).astype(F32) + hit = cum >= target[:, None] + idx = np.argmax(hit, axis=1) + idx[~hit.any(axis=1)] = V - 1 # reference fallback + sampled[c0:c1] = idx + prob = e / Z[:, None] + with np.errstate(divide="ignore", invalid="ignore"): + h = np.where(prob > 0, prob * np.log(prob), F32(0)) + entropy[c0:c1] = -np.cumsum(h, axis=1, dtype=F32)[:, -1] + + order = np.argsort(entropy, kind="stable") # std::sort: ties unspecified + cumE = np.cumsum(entropy[order].astype(np.float64)) + ok = (cumE - entropy[order]) <= np.float64(F32(p.entropy_bound)) + accepted = np.zeros(C, bool) + accepted[order[ok]] = True + + canvas_next = np.where(accepted, sampled, renoise).astype(np.int32) + argmax_i32 = amax.astype(np.int32) + entropy_sum = np.cumsum(entropy, dtype=F32)[-1] # sequential f32 like the C++ + + self.held = self.held + 1 if np.array_equal(self._prev_argmax, argmax_i32) else 0 + confident = (entropy_sum / F32(C)) < F32(p.confidence_threshold) + if self.held >= p.stability_threshold and confident: + self.finished = True + self._prev_argmax = argmax_i32 + self.prev_temp_inv = temp_inv + self.argmax_canvas = argmax_i32 + self.canvas = canvas_next + self.step_idx += 1 + if self.step_idx >= self.S: + self.finished = True + + return StepResult(step_idx, cur_step, float(t), entropy, argmax_i32, + sampled.astype(np.int32), accepted, canvas_next, + u, renoise.astype(np.int32), self.held, self.finished, + float(entropy_sum / F32(C))) diff --git a/temp/diffgemma_s2/compare_eb.py b/temp/diffgemma_s2/compare_eb.py new file mode 100644 index 0000000..7400555 --- /dev/null +++ b/temp/diffgemma_s2/compare_eb.py @@ -0,0 +1,121 @@ +"""compare_eb.py — drive the SK EntropyBoundSampler with the same per-step +logits a reference run consumed and diff every decision field against the +reference's .dec records (eb_ref_harness locally, or the instrumented +llama.cpp diffusion.cpp on amelia — identical record layout). + + python3 compare_eb.py --logits-dir d --dec-dir d2 --seed s --C 256 \ + --S 16 --n-vocab 262144 [--params t_min t_max bound stab conf] +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +from SuperKittens.models.gemma.diffusion.sampler import EBParams, EntropyBoundSampler # noqa: E402 + + +def read_dec(path: Path, C: int) -> dict: + b = path.read_bytes() + off = 0 + + def take(dtype, n): + nonlocal off + a = np.frombuffer(b, dtype=dtype, count=n, offset=off) + off += a.nbytes + return a + + d = {} + d["step_idx"] = int(take(np.int32, 1)[0]) + d["cur_step"] = int(take(np.int32, 1)[0]) + d["t"] = float(take(np.float32, 1)[0]) + d["canvas_in"] = take(np.int32, C) + d["u"] = take(np.float32, C) + d["renoise"] = take(np.int32, C) + d["entropy"] = take(np.float32, C) + d["argmax"] = take(np.int32, C) + d["denoiser"] = take(np.int32, C) + d["accepted"] = take(np.uint8, C).astype(bool) + d["canvas_next"] = take(np.int32, C) + d["held"] = int(take(np.int32, 1)[0]) + d["finished"] = bool(take(np.uint8, 1)[0]) + d["entropy_sum"] = float(take(np.float32, 1)[0]) + return d + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--logits-dir", required=True) + ap.add_argument("--dec-dir", required=True) + ap.add_argument("--seed", type=int, required=True) + ap.add_argument("--C", type=int, default=256) + ap.add_argument("--S", type=int, required=True) + ap.add_argument("--n-vocab", type=int, default=262144) + ap.add_argument("--params", nargs=5, type=float, default=None, + metavar=("TMIN", "TMAX", "BOUND", "STAB", "CONF")) + args = ap.parse_args() + + p = EBParams(max_steps=args.S, seed=args.seed) + if args.params: + p.t_min, p.t_max, p.entropy_bound = args.params[0], args.params[1], args.params[2] + p.stability_threshold = int(args.params[3]) + p.confidence_threshold = args.params[4] + + smp = EntropyBoundSampler(p, args.n_vocab, args.C) + dec_dir = Path(args.dec_dir) + ldir = Path(args.logits_dir) + + n_steps = len(sorted(dec_dir.glob("step_*.dec"))) + total = mism = 0 + fields = ["canvas_in", "u", "renoise", "entropy", "argmax", "denoiser", + "accepted", "canvas_next", "held", "finished"] + per_field = {f: 0 for f in fields} + + for s in range(n_steps): + ref = read_dec(dec_dir / f"step_{s:03d}.dec", args.C) + canvas_before = smp.canvas.copy() + logits = np.fromfile(ldir / f"step_{s:03d}.f32", dtype=np.float32).reshape( + args.C, args.n_vocab) + r = smp.step(logits) + mine = {"canvas_in": canvas_before, "u": r.u, "renoise": r.renoise, + "entropy": r.entropy, "argmax": r.argmax, "denoiser": r.sampled, + "accepted": r.accepted, "canvas_next": r.canvas_next, + "held": r.held, "finished": r.finished} + line = [f"step {s:3d} t={r.t:.4f}"] + for f in fields: + a, b = mine[f], ref[f] + if f == "entropy": + d = float(np.abs(a - b).max()) + ok = d < 5e-4 + line.append(f"H~{d:.2e}") + elif isinstance(a, (int, bool)): + ok = a == b + else: + ok = np.array_equal(a, b) + if not ok: + n_bad = int((np.asarray(a) != np.asarray(b)).sum()) + line.append(f"{f}:{n_bad}bad") + total += 1 + if not ok: + mism += 1 + per_field[f] += 1 + naccept = int(r.accepted.sum()) + line.append(f"acc={naccept} held={r.held} Hbar={r.entropy_mean:.4f}" + f" fin={r.finished}") + print(" ".join(line)) + if r.finished: + break + + print(f"\nfields compared: {total}, mismatched: {mism}") + for f, n in per_field.items(): + if n: + print(f" {f}: {n} steps mismatched") + print("TOKEN-IDENTICAL" if mism == 0 else "MISMATCH") + return 0 if mism == 0 else 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/temp/diffgemma_s2/eb_ref_harness.cpp b/temp/diffgemma_s2/eb_ref_harness.cpp new file mode 100644 index 0000000..2e35a8c --- /dev/null +++ b/temp/diffgemma_s2/eb_ref_harness.cpp @@ -0,0 +1,194 @@ +// eb_ref_harness.cpp — diffusion_generate_entropy_bound lifted VERBATIM from +// llama.cpp PR #24423 examples/diffusion/diffusion.cpp @ c84e85af, with +// llama_decode replaced by per-step logits files. Drives the same logits into +// the C++ decision path so the Python sampler can be checked token-for-token +// in isolation from any model/forward noise. +// +// ./eb_ref_harness +// [t_min t_max entropy_bound stability confidence] +// +// logits_dir/step_%03d.f32 : [C, n_vocab] f32 (step_idx-indexed; the harness +// stops consuming at adaptive stop, like the ref) +// out_dir/step_%03d.dec : decision record, parsed by compare_eb.py +// i32 step_idx, i32 cur_step, f32 t, +// i32 canvas_in[C], f32 u[C], i32 renoise[C], f32 entropy[C], +// i32 argmax[C], i32 denoiser[C], u8 accepted[C], i32 canvas_next[C], +// i32 held, u8 finished, f32 entropy_sum +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +typedef int32_t llama_token; + +struct eb_params { + int32_t S; + float t_min = 0.4f; + float t_max = 0.8f; + float entropy_bound = 0.1f; + int32_t stability_threshold = 1; + float confidence_threshold = 0.005f; + int32_t seed = 0; +}; + +int main(int argc, char** argv) { + if (argc < 7) { fprintf(stderr, "usage: see header\n"); return 1; } + const std::string ldir = argv[1]; + const std::string odir = argv[2]; + eb_params params; + params.seed = atoi(argv[3]); + const int32_t C = atoi(argv[4]); + params.S = atoi(argv[5]); + const int32_t n_vocab = atoi(argv[6]); + if (argc >= 12) { + params.t_min = (float) atof(argv[7]); + params.t_max = (float) atof(argv[8]); + params.entropy_bound = (float) atof(argv[9]); + params.stability_threshold = atoi(argv[10]); + params.confidence_threshold = (float) atof(argv[11]); + } + const int32_t S = params.S; + + // ---- verbatim reference body below (file-fed logits) ------------------- + std::mt19937 rng(params.seed); + std::uniform_real_distribution uni01(0.0f, 1.0f); + std::uniform_int_distribution vocab_dist(0, n_vocab - 1); + + std::vector current_canvas(C); + for (int32_t i = 0; i < C; i++) { + current_canvas[i] = vocab_dist(rng); + } + + std::vector argmax_canvas(C, 0); + std::vector prev_argmax(C, -1); + std::vector entropy(C); + std::vector denoiser(C); + std::vector order(C); + std::vector u(C); + std::vector renoise(C); + + const unsigned hw = std::thread::hardware_concurrency(); + const unsigned nth = std::max(1u, std::min(hw ? hw : 1u, 32u)); + + std::vector logits((size_t) C * n_vocab); + + float prev_temp_inv = 1.0f; + int held = 0; + bool finished = false; + (void) prev_temp_inv; + + for (int32_t cur_step = S; cur_step >= 1 && !finished; --cur_step) { + const int32_t step_idx = S - cur_step; + const float t = params.t_min + (params.t_max - params.t_min) * ((float) cur_step / (float) S); + const float temp_inv = 1.0f / t; + + // forward stand-in: read this step's logits + { + char path[1024]; + snprintf(path, sizeof(path), "%s/step_%03d.f32", ldir.c_str(), step_idx); + FILE* f = fopen(path, "rb"); + if (!f) { fprintf(stderr, "missing %s\n", path); return 1; } + if (fread(logits.data(), 4, logits.size(), f) != logits.size()) { + fprintf(stderr, "short read %s\n", path); return 1; + } + fclose(f); + } + std::vector canvas_in = current_canvas; + + for (int32_t pos = 0; pos < C; pos++) { + u[pos] = uni01(rng); + renoise[pos] = vocab_dist(rng); + } + + auto worker = [&](int32_t p0, int32_t p1) { + for (int32_t pos = p0; pos < p1; pos++) { + const float* row = logits.data() + (size_t) pos * n_vocab; + float m = -INFINITY; int32_t amax = 0; + for (int32_t v = 0; v < n_vocab; v++) { + const float z = row[v] * temp_inv; + if (z > m) { m = z; amax = v; } + } + float Z = 0.0f; + for (int32_t v = 0; v < n_vocab; v++) { + Z += expf(row[v] * temp_inv - m); + } + const float target = u[pos] * Z; + float cum = 0.0f, H = 0.0f; + int32_t sampled = n_vocab - 1; bool picked = false; + for (int32_t v = 0; v < n_vocab; v++) { + const float e = expf(row[v] * temp_inv - m); + const float p = e / Z; + if (p > 0.0f) { H -= p * logf(p); } + cum += e; + if (!picked && cum >= target) { sampled = v; picked = true; } + } + entropy[pos] = H; + argmax_canvas[pos] = amax; + denoiser[pos] = sampled; + } + }; + { + std::vector pool; + const int32_t chunk = (C + (int32_t) nth - 1) / (int32_t) nth; + for (unsigned ti = 0; ti < nth; ti++) { + const int32_t p0 = (int32_t) ti * chunk; + const int32_t p1 = std::min(p0 + chunk, C); + if (p0 < p1) { pool.emplace_back(worker, p0, p1); } + } + for (auto& th : pool) { th.join(); } + } + + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](int32_t a, int32_t b) { return entropy[a] < entropy[b]; }); + std::vector accepted(C, 0); + double cumE = 0.0; + for (int32_t k = 0; k < C; k++) { + const int32_t pos = order[k]; + cumE += entropy[pos]; + if (cumE - entropy[pos] <= params.entropy_bound) { accepted[pos] = 1; } + } + + float entropy_sum = 0.0f; + for (int32_t pos = 0; pos < C; pos++) { + current_canvas[pos] = accepted[pos] ? denoiser[pos] : renoise[pos]; + entropy_sum += entropy[pos]; + } + + held = (prev_argmax == argmax_canvas) ? held + 1 : 0; + const bool confident = (entropy_sum / (float) C) < params.confidence_threshold; + if (held >= params.stability_threshold && confident) { finished = true; } + prev_argmax = argmax_canvas; + prev_temp_inv = temp_inv; + + // ---- decision dump -------------------------------------------------- + { + char path[1024]; + snprintf(path, sizeof(path), "%s/step_%03d.dec", odir.c_str(), step_idx); + FILE* f = fopen(path, "wb"); + if (!f) { perror("fopen dec"); return 1; } + uint8_t fin = finished ? 1 : 0; + fwrite(&step_idx, 4, 1, f); fwrite(&cur_step, 4, 1, f); fwrite(&t, 4, 1, f); + fwrite(canvas_in.data(), 4, C, f); + fwrite(u.data(), 4, C, f); + fwrite(renoise.data(), 4, C, f); + fwrite(entropy.data(), 4, C, f); + fwrite(argmax_canvas.data(), 4, C, f); + fwrite(denoiser.data(), 4, C, f); + std::vector acc8(accepted.begin(), accepted.end()); + fwrite(acc8.data(), 1, C, f); + fwrite(current_canvas.data(), 4, C, f); + fwrite(&held, 4, 1, f); + fwrite(&fin, 1, 1, f); + fwrite(&entropy_sum, 4, 1, f); + fclose(f); + } + } + return 0; +} diff --git a/temp/diffgemma_s2/rng_dump.cpp b/temp/diffgemma_s2/rng_dump.cpp new file mode 100644 index 0000000..ceee885 --- /dev/null +++ b/temp/diffgemma_s2/rng_dump.cpp @@ -0,0 +1,49 @@ +// rng_dump.cpp — emit the exact RNG stream the reference EB sampler consumes +// (std::mt19937 + libc++ uniform_int/uniform_real, same construction & call +// order as diffusion_generate_entropy_bound) so the Python mirror can be +// verified bit-for-bit. +// +// ./rng_dump +// +// layout (little-endian): +// int32 canvas_init[C] +// per step s in 0..S-1: float32 u[C], int32 renoise[C] +// (u/renoise interleave per position inside the generator, matching the +// reference's "pre-draw single-threaded" loop) +#include +#include +#include +#include +#include + +int main(int argc, char** argv) { + if (argc != 6) { fprintf(stderr, "usage: %s seed C S n_vocab out\n", argv[0]); return 1; } + const int32_t seed = atoi(argv[1]); + const int32_t C = atoi(argv[2]); + const int32_t S = atoi(argv[3]); + const int32_t n_vocab = atoi(argv[4]); + + std::mt19937 rng(seed); + std::uniform_real_distribution uni01(0.0f, 1.0f); + std::uniform_int_distribution vocab_dist(0, n_vocab - 1); + + FILE* f = fopen(argv[5], "wb"); + if (!f) { perror("fopen"); return 1; } + + std::vector canvas(C); + for (int32_t i = 0; i < C; i++) canvas[i] = vocab_dist(rng); + fwrite(canvas.data(), 4, C, f); + + std::vector u(C); + std::vector renoise(C); + for (int32_t s = 0; s < S; s++) { + for (int32_t pos = 0; pos < C; pos++) { + u[pos] = uni01(rng); + renoise[pos] = vocab_dist(rng); + } + fwrite(u.data(), 4, C, f); + fwrite(renoise.data(), 4, C, f); + } + fclose(f); + return 0; +} From 109d131bccee524f77ed8120401da99c8a531e3c Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 23:24:42 -0400 Subject: [PATCH 18/31] preserve: diffgemma stage2 WIP (agent session-limit recovery; SC subgraph + e2e generate driver in progress, ungated) --- .../models/gemma/diffusion/forward_metal.py | 84 ++++++++- .../models/gemma/diffusion/generate.py | 167 ++++++++++++++++++ .../models/gemma/diffusion/graph_ref.py | 42 ++++- SuperKittens/models/gemma/diffusion/runner.py | 16 +- 4 files changed, 297 insertions(+), 12 deletions(-) create mode 100644 SuperKittens/models/gemma/diffusion/generate.py diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index c0225da..6debffb 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -26,7 +26,7 @@ from .config import DiffusionGemmaConfig from .gguf_io import GGUFFile from .graph_ref import (F32, Weights, build_mask, embed_tokens, gelu_tanh, - moe_route, rms_norm, rope_neox) + moe_route, rms_norm, rope_neox, softmax) _SK_ROOT = Path(__file__).resolve().parents[3] _KERNEL_SOURCES = [ @@ -204,13 +204,82 @@ def run(self): class DiffusionGemmaMetal: """Stage-1 unified zero-SC forward. forward(ids, P) -> canvas logits f32.""" - def __init__(self, gguf_path: str, cfg: DiffusionGemmaConfig): + def __init__(self, gguf_path: str, cfg: DiffusionGemmaConfig, + sc_embt_path: str | None = None): self.gg = GGUFFile(gguf_path) self.cfg = cfg self.ctx = MetalCtx() self.wb = WeightBufs(self.ctx, self.gg) self.w = Weights(self.gg) # F32 sidecars (norms, scales, router) self.dump = None # optional (name, il, arr) tap + # self-conditioning soft-embed: transposed dequantized embed + # [d_model, vocab] f16 on disk (make_embt.py), streamed per chunk + # through a persistent slot like every other weight + self.sc_embt_path = sc_embt_path + self._sc_bufs = None # (probs_buf, chunk_buf, out_buf, mm) + self.sc_chunk_rows = 352 # 8 chunks x 352 = 2816; 184.5 MB slot + + # -- self-conditioning ------------------------------------------------------- + + def _sc_soft_embed(self, probs16: np.ndarray) -> np.ndarray: + """probs16 fp16 [C, V] @ embed [V, d] -> f32 [C, d] via the streamed + transposed-embed GEMM (dg_gemm_qkt_f32: f16 inputs, f32 C — matches + the reference's f16 sc_embT matmul, f32-accumulated).""" + cfg = self.cfg + C, V = probs16.shape + d = cfg.d_model + rows = self.sc_chunk_rows + assert d % rows == 0 + if self._sc_bufs is None: + f = open(self.sc_embt_path, "rb") + mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ) + assert len(mm) == d * V * 2, "embT size mismatch (expect [d, V] f16)" + self._sc_bufs = (self.ctx.buf_empty(C * V * 2), + self.ctx.buf_empty(rows * V * 2), + self.ctx.buf_empty(C * d * 4), mm) + p_buf, w_buf, o_buf, mm = self._sc_bufs + p_buf.contents().as_buffer(C * V * 2)[:] = probs16.tobytes() + for r0 in range(0, d, rows): + nbytes = rows * V * 2 + w_buf.contents().as_buffer(nbytes)[:] = mm[r0 * V * 2:r0 * V * 2 + nbytes] + b = GemmBatch(self.ctx) + b.gemm("dg_gemm_qkt_f32", p_buf, 0, w_buf, 0, o_buf, r0 * 4, + M=C, N=rows, K=V, ldc=d) + b.run() + try: + pg = mmap.PAGESIZE + a = (r0 * V * 2) // pg * pg + mm.madvise(mmap.MADV_DONTNEED, a, + (nbytes + (r0 * V * 2) - a + pg - 1) // pg * pg) + except (AttributeError, ValueError, OSError): + pass + return self.ctx.read(o_buf, np.float32, (C, d)) + + def _sc_signal(self, sc_logits: np.ndarray, sc_temp_inv: float, + sc_use: float) -> np.ndarray: + """PR dg_canvas_embed SC subgraph -> sc_sig f32 [C, d_model].""" + cfg, w = self.cfg, self.w + C, V = sc_logits.shape + # softmax(prev raw logits / prev t), fp16 on the wire like the + # reference (ggml converts the f32 probs to the f16 vec_dot type) + probs16 = np.empty((C, V), np.float16) + for c0 in range(0, C, 32): + c1 = min(c0 + 32, C) + probs16[c0:c1] = softmax(sc_logits[c0:c1] * F32(sc_temp_inv)).astype(np.float16) + if self.sc_embt_path is not None: + soft = self._sc_soft_embed(probs16) + else: # oracle-style host fallback (slow: full embed dequant per step) + soft = np.zeros((C, cfg.d_model), F32) + pf = probs16.astype(F32) + for v0 in range(0, V, 16384): + v1 = min(v0 + 16384, V) + soft += pf[:, v0:v1] @ self.w.dq("token_embd.weight", rows=slice(v0, v1)) + soft = soft * F32(np.sqrt(F32(cfg.d_model))) + normed = rms_norm(soft, cfg.eps, w.f32("self_cond_pre_norm.weight")) + g = gelu_tanh(self._gemm_f32("self_cond_gate.weight", normed, cfg.n_ff)) + u = self._gemm_f32("self_cond_up.weight", normed, cfg.n_ff) + sig = self._gemm_f32("self_cond_down.weight", g * u, cfg.d_model) + return (sig * F32(sc_use)).astype(F32) # -- helpers --------------------------------------------------------------- @@ -331,7 +400,9 @@ def _moe(self, il: int, attn_out: np.ndarray, e_in: np.ndarray) -> np.ndarray: # -- forward --------------------------------------------------------------- - def forward(self, ids: np.ndarray, P: int) -> np.ndarray: + def forward(self, ids: np.ndarray, P: int, sc_logits: np.ndarray | None = None, + sc_temp_inv: float = 1.0, sc_use: float = 1.0) -> np.ndarray: + """sc_logits=None -> the Stage-1-validated zero-SC unified forward.""" cfg, w = self.cfg, self.w ids = np.asarray(ids) N = len(ids) @@ -339,7 +410,12 @@ def forward(self, ids: np.ndarray, P: int) -> np.ndarray: pos = np.arange(N, dtype=np.int64) dmp = self.dump or (lambda name, il, arr: None) - x = embed_tokens(w, cfg, ids, P) + sc_sig = None + if sc_logits is not None: + with objc.autorelease_pool(): + sc_sig = self._sc_signal(sc_logits, sc_temp_inv, sc_use) + dmp("sc_sig", -1, sc_sig) + x = embed_tokens(w, cfg, ids, P, sc_sig=sc_sig) dmp("inp_region", -1, x) rope_ff = w.f32("rope_freqs.weight") diff --git a/SuperKittens/models/gemma/diffusion/generate.py b/SuperKittens/models/gemma/diffusion/generate.py new file mode 100644 index 0000000..aedca80 --- /dev/null +++ b/SuperKittens/models/gemma/diffusion/generate.py @@ -0,0 +1,167 @@ +# pyright: reportMissingImports=false +"""generate.py — DiffusionGemma end-to-end block generation (Stage 2). + +One denoising block: EntropyBound sampler (sampler.py, reference +decision-mirror) driving the unified [prompt|canvas] forward with +self-conditioning. Per-step telemetry to stdout + a JSONL log; final output = +the argmax canvas, trimmed at the first end-of-generation token / repetition +loop like the reference CLI. + + python -m SuperKittens.models.gemma.diffusion.generate \ + --gguf model.gguf --prompt-ids p.i32 --out-dir gen/ \ + [--steps 16 --seed 1234 --mode gpu --sc-embt dg_embT_f16.bin --no-sc] +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np + +from .config import config_from_gguf +from .gguf_io import GGUFFile +from .sampler import EBParams, EntropyBoundSampler + +F32 = np.float32 + + +def detokenize(meta: dict, ids) -> str: + """Minimal SPM detok (ggml llama-style vocab): ▁ -> space, <0xXX> bytes.""" + tokens = meta["tokenizer.ggml.tokens"] + out = bytearray() + for tid in ids: + piece = tokens[int(tid)] + if len(piece) == 6 and piece.startswith("<0x") and piece.endswith(">"): + out += bytes([int(piece[3:5], 16)]) + else: + out += piece.replace("▁", " ").encode("utf-8") + return out.decode("utf-8", "replace") + + +def eog_ids(meta: dict) -> set[int]: + ids = set() + for k in ("tokenizer.ggml.eos_token_id", "tokenizer.ggml.eot_token_id"): + if k in meta: + ids.add(int(meta[k])) + return ids + + +def trim_canvas(canvas: np.ndarray, eog: set[int]) -> int: + """Reference CLI trim: cut at the first EOG token, else at the onset of a + stride-1/2 repetition loop (>= 6 reps).""" + n = len(canvas) + cut = n + for i in range(n): + if int(canvas[i]) in eog: + cut = i + break + for i in range(cut - 1): + for stride in (1, 2): + reps = 0 + j = i + while j + stride < n and canvas[j] == canvas[j + stride]: + reps += 1 + j += stride + if reps >= 6: + return i + return cut + + +def run_block(model, cfg, prompt_ids: np.ndarray, params: EBParams, + use_sc: bool, log, mode: str = "gpu"): + """One denoising block; returns (argmax_canvas, steps_run, timings).""" + C = cfg.canvas_length + P = len(prompt_ids) + smp = EntropyBoundSampler(params, cfg.vocab_size, C) + prev_logits = None + fw_s = smp_s = 0.0 + step = None + while not smp.finished: + ids = np.concatenate([prompt_ids, smp.canvas.astype(np.int32)]) + t0 = time.time() + if use_sc: + sc = prev_logits if prev_logits is not None else np.zeros((C, cfg.vocab_size), F32) + sc_use = 0.0 if smp.step_idx == 0 else 1.0 + logits = model.forward(ids, P, sc_logits=sc, + sc_temp_inv=float(smp.prev_temp_inv), sc_use=sc_use) + else: + logits = model.forward(ids, P) + t1 = time.time() + if not np.isfinite(logits).all(): + raise RuntimeError(f"non-finite logits at step {smp.step_idx}") + step = smp.step(np.ascontiguousarray(logits, dtype=F32)) + t2 = time.time() + fw_s += t1 - t0 + smp_s += t2 - t1 + prev_logits = logits + rec = {"step": step.step_idx, "t": round(step.t, 4), + "accepted": int(step.accepted.sum()), "held": step.held, + "H_mean": round(step.entropy_mean, 5), "finished": step.finished, + "fw_s": round(t1 - t0, 1), "smp_s": round(t2 - t1, 1)} + log(rec) + return smp.argmax_canvas, (step.step_idx + 1 if step else 0), (fw_s, smp_s) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--gguf", required=True) + ap.add_argument("--prompt-ids", required=True, help="chat-templated prompt, raw i32") + ap.add_argument("--out-dir", required=True) + ap.add_argument("--mode", choices=["gpu", "cpu"], default="gpu") + ap.add_argument("--steps", type=int, default=16) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--sc-embt", default=None, help="dg_embT_f16.bin (make_embt.py)") + ap.add_argument("--no-sc", action="store_true", help="zero-SC forward (Stage-1 config)") + args = ap.parse_args(argv) + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + prompt = np.fromfile(args.prompt_ids, dtype=np.int32) + gg = GGUFFile(args.gguf) + cfg = config_from_gguf(gg.meta) + params = EBParams(max_steps=args.steps, seed=args.seed) + + if args.mode == "cpu": + from .graph_ref import forward_cpu + + class _CpuModel: + def forward(self, ids, P, **sc): + return forward_cpu(gg, cfg, ids, P, **sc) + model = _CpuModel() + else: + from .forward_metal import DiffusionGemmaMetal + model = DiffusionGemmaMetal(args.gguf, cfg, sc_embt_path=args.sc_embt) + + jl = open(out / "steps.jsonl", "w") + + def log(rec): + print(json.dumps(rec), flush=True) + jl.write(json.dumps(rec) + "\n") + jl.flush() + + t0 = time.time() + canvas, steps, (fw_s, smp_s) = run_block(model, cfg, prompt, params, + use_sc=not args.no_sc, log=log, + mode=args.mode) + wall = time.time() - t0 + + canvas.astype(np.int32).tofile(out / "canvas.i32") + cut = trim_canvas(canvas, eog_ids(gg.meta)) + text = detokenize(gg.meta, canvas[:cut]) + (out / "output.txt").write_text(text) + n_mask = int((canvas == cfg.mask_token_id).sum()) + summary = {"steps": steps, "wall_s": round(wall, 1), "fw_s": round(fw_s, 1), + "smp_s": round(smp_s, 1), "trim": cut, "mask_tokens_in_canvas": n_mask, + "seed": args.seed, "S": args.steps, "sc": not args.no_sc} + (out / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps(summary)) + print("---- output ----") + print(text) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/SuperKittens/models/gemma/diffusion/graph_ref.py b/SuperKittens/models/gemma/diffusion/graph_ref.py index 1b6748e..7c65c3b 100644 --- a/SuperKittens/models/gemma/diffusion/graph_ref.py +++ b/SuperKittens/models/gemma/diffusion/graph_ref.py @@ -108,13 +108,39 @@ def dq(self, name: str, rows=None) -> np.ndarray: # -- region-aware unified forward (CPU f32 reference) ------------------------- -def embed_tokens(w: Weights, cfg: DiffusionGemmaConfig, ids: np.ndarray, P: int) -> np.ndarray: +def embed_tokens(w: Weights, cfg: DiffusionGemmaConfig, ids: np.ndarray, P: int, + sc_sig: np.ndarray | None = None) -> np.ndarray: + """Region embedding. Canvas rows: rmsnorm_noscale(embed*sqrt(d) [+ sc_sig]); + sc_sig is the (already sc_use-gated) self-conditioning signal.""" x = w.dq("token_embd.weight", rows=np.asarray(ids, np.int64)) x = x * F32(np.sqrt(F32(cfg.d_model))) - x[P:] = rms_norm(x[P:], cfg.eps) # canvas rows: rmsnorm no-scale (zero-SC) + canvas = x[P:] if sc_sig is None else (x[P:] + sc_sig).astype(F32) + x[P:] = rms_norm(canvas, cfg.eps) return x.astype(F32) +def sc_signal(w: Weights, cfg: DiffusionGemmaConfig, sc_logits: np.ndarray, + sc_temp_inv: float, sc_use: float, soft: np.ndarray | None = None, + _vchunk: int = 16384) -> np.ndarray: + """Self-conditioning gated MLP (PR dg_canvas_embed): softmax(prev raw + logits / prev t) -> soft-embedding -> rms*pre_norm -> down(gelu(gate)*up), + scaled by the runtime {0,1} sc_use gate. `soft` lets the Metal driver + inject its GPU soft-embedding; the oracle computes it chunked f32.""" + C = sc_logits.shape[0] + if soft is None: + probs = softmax(sc_logits * F32(sc_temp_inv)) + soft = np.zeros((C, cfg.d_model), dtype=F32) + for v0 in range(0, cfg.vocab_size, _vchunk): + v1 = min(v0 + _vchunk, cfg.vocab_size) + soft += probs[:, v0:v1] @ w.dq("token_embd.weight", rows=slice(v0, v1)) + soft = soft * F32(np.sqrt(F32(cfg.d_model))) + normed = rms_norm(soft, cfg.eps, w.f32("self_cond_pre_norm.weight")) + g = gelu_tanh(normed @ w.dq("self_cond_gate.weight").T) + u = normed @ w.dq("self_cond_up.weight").T + sig = (g * u) @ w.dq("self_cond_down.weight").T + return (sig * F32(sc_use)).astype(F32) + + def moe_route(w: Weights, cfg: DiffusionGemmaConfig, attn_out: np.ndarray, il: int): """Router (operates on the UNNORMED residual): rms_noscale -> /sqrt(d) -> * gate_inp scale -> logits -> softmax -> top-8 -> renorm weights.""" @@ -130,15 +156,21 @@ def moe_route(w: Weights, cfg: DiffusionGemmaConfig, attn_out: np.ndarray, il: i def forward_cpu(gg: GGUFFile, cfg: DiffusionGemmaConfig, ids: np.ndarray, - P: int, dump=None) -> np.ndarray: - """Unified [prompt|canvas] zero-SC forward; returns canvas logits f32 [C, V].""" + P: int, dump=None, sc_logits: np.ndarray | None = None, + sc_temp_inv: float = 1.0, sc_use: float = 1.0) -> np.ndarray: + """Unified [prompt|canvas] forward; returns canvas logits f32 [C, V]. + sc_logits=None -> the Stage-1-validated zero-SC forward.""" w = Weights(gg) N = len(ids) C = N - P pos = np.arange(N, dtype=np.int64) dmp = dump or (lambda name, il, arr: None) - x = embed_tokens(w, cfg, np.asarray(ids), P) + sc_sig = None + if sc_logits is not None: + sc_sig = sc_signal(w, cfg, sc_logits, sc_temp_inv, sc_use) + dmp("sc_sig", -1, sc_sig) + x = embed_tokens(w, cfg, np.asarray(ids), P, sc_sig=sc_sig) dmp("inp_region", -1, x) rope_ff = w.f32("rope_freqs.weight") diff --git a/SuperKittens/models/gemma/diffusion/runner.py b/SuperKittens/models/gemma/diffusion/runner.py index a92eaca..6d2f6eb 100644 --- a/SuperKittens/models/gemma/diffusion/runner.py +++ b/SuperKittens/models/gemma/diffusion/runner.py @@ -49,6 +49,10 @@ def main(argv=None) -> int: ap.add_argument("--layers", type=int, default=0, help="truncate to first N layers (debug)") ap.add_argument("--dump-dir", default=None) ap.add_argument("--dump-names", default="l_out") + ap.add_argument("--sc-logits", default=None, + help="prev-step raw canvas logits f32 [C,V] (enables SC, use_sc=1)") + ap.add_argument("--sc-temp-inv", type=float, default=1.0) + ap.add_argument("--sc-embt", default=None, help="dg_embT_f16.bin for the GPU SC path") args = ap.parse_args(argv) prompt = np.fromfile(args.prompt_ids, dtype=np.int32) @@ -66,15 +70,21 @@ def main(argv=None) -> int: cfg.n_layers = args.layers dump = make_dump(args.dump_dir, set(filter(None, args.dump_names.split(",")))) + sc = {} + if args.sc_logits: + sc_arr = np.fromfile(args.sc_logits, dtype=np.float32).reshape( + cfg.canvas_length, cfg.vocab_size) + sc = {"sc_logits": sc_arr, "sc_temp_inv": args.sc_temp_inv, "sc_use": 1.0} + t0 = time.time() if args.mode == "cpu": from .graph_ref import forward_cpu - logits = forward_cpu(gg, cfg, ids, P, dump=dump) + logits = forward_cpu(gg, cfg, ids, P, dump=dump, **sc) else: from .forward_metal import DiffusionGemmaMetal - m = DiffusionGemmaMetal(args.gguf, cfg) + m = DiffusionGemmaMetal(args.gguf, cfg, sc_embt_path=args.sc_embt) m.dump = dump - logits = m.forward(ids, P) + logits = m.forward(ids, P, **sc) dt = time.time() - t0 logits.astype(np.float32).tofile(args.out) From d249be99d5fbd175aadb7c6415ea5171f8f2560a Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 23:29:21 -0400 Subject: [PATCH 19/31] best.md: Phi-4-reasoning 14B row (10.49 tok/s amelia, config-only port) --- best.md | 1 + 1 file changed, 1 insertion(+) diff --git a/best.md b/best.md index 6485be5..2ac38ea 100644 --- a/best.md +++ b/best.md @@ -31,6 +31,7 @@ vs llama.cpp on a clean-`main` regression sweep. | nemotron-Nano-8B (Llama-3.1) | **21.58** | derek | yes | interleaved/NORM RoPE for Llama GGUFs fixed the degeneration ([#79]); shared dense core | | gemma4-E2B (PLE, KV-share) | **~21** | derek | yes | Q8 body+lmhead ([#69]/[#72]) | | Mistral-7B-Instruct-v0.3 | **19.93** | amelia | yes | config-only over DenseDecoder + interleaved RoPE ([#84]) | +| Phi-4-reasoning 14B (phi3 arch) | **10.49** | amelia | yes | config-only over DenseDecoder + one-time fused-tensor GGUF repack; amelia-with-colima number — expect ≥Qwen3-14B on lexie (local main, June 10) | | gemma4-E4B (own ckpt, 42L) | **9.86** | — | yes | PLE-table + embed Q8 fixed the 1.44→9.86 paging cliff ([#74]) | | gemma4-unified-12B (distinct arch) | **8.2–8.4** | derek | yes | Q4_K body (q4k/q6k_matvec_bf16) fits under the ~12 GB Metal wired limit + fp16-subnormal embed-dequant fix ([#78]) | From 9946f51e010d985dc7d451c8a6d1f8e1c43d301a Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 10 Jun 2026 23:33:22 -0400 Subject: [PATCH 20/31] diffgemma stage2: commit recovered lab tools (make_embt, instrument_diffusion); generate.py skips SC compute at step 0 (sc_use=0 is bit-identical to zero-SC, saves the 1.5 GB embT stream) --- .../models/gemma/diffusion/generate.py | 10 +- temp/diffgemma_s2/instrument_diffusion.py | 99 +++++++++++++++++++ temp/diffgemma_s2/make_embt.py | 43 ++++++++ 3 files changed, 147 insertions(+), 5 deletions(-) create mode 100644 temp/diffgemma_s2/instrument_diffusion.py create mode 100644 temp/diffgemma_s2/make_embt.py diff --git a/SuperKittens/models/gemma/diffusion/generate.py b/SuperKittens/models/gemma/diffusion/generate.py index aedca80..b839fd4 100644 --- a/SuperKittens/models/gemma/diffusion/generate.py +++ b/SuperKittens/models/gemma/diffusion/generate.py @@ -82,11 +82,11 @@ def run_block(model, cfg, prompt_ids: np.ndarray, params: EBParams, while not smp.finished: ids = np.concatenate([prompt_ids, smp.canvas.astype(np.int32)]) t0 = time.time() - if use_sc: - sc = prev_logits if prev_logits is not None else np.zeros((C, cfg.vocab_size), F32) - sc_use = 0.0 if smp.step_idx == 0 else 1.0 - logits = model.forward(ids, P, sc_logits=sc, - sc_temp_inv=float(smp.prev_temp_inv), sc_use=sc_use) + if use_sc and prev_logits is not None: + # step 0 (sc_use=0 in the reference) is bit-identical to zero-SC: + # sig*0 == 0 and x+0 == x — skip the 1.5 GB embT stream entirely + logits = model.forward(ids, P, sc_logits=prev_logits, + sc_temp_inv=float(smp.prev_temp_inv), sc_use=1.0) else: logits = model.forward(ids, P) t1 = time.time() diff --git a/temp/diffgemma_s2/instrument_diffusion.py b/temp/diffgemma_s2/instrument_diffusion.py new file mode 100644 index 0000000..413de1c --- /dev/null +++ b/temp/diffgemma_s2/instrument_diffusion.py @@ -0,0 +1,99 @@ +"""instrument_diffusion.py — add env-driven per-step dumps to llama.cpp PR +#24423 examples/diffusion/diffusion.cpp (diffusion_generate_entropy_bound). + +DG_EB_DUMP= : per step write step_%03d.f32 (canvas logits [C, V]) + and step_%03d.dec (decision record, eb_ref_harness + layout) + header.bin (n_input,C,S,seed, eb params, + prompt ids) at step 0. +DG_EB_DUMP_THROTTLE=1 : block until the consumer deletes step-2's logits + (keeps <=2 x 268 MB in flight on a tight disk). + +Idempotent (skips if already instrumented). Usage: + python3 instrument_diffusion.py +""" +import sys +from pathlib import Path + +A_INCLUDE = "#include \n" +INS_INCLUDE = "#include \n#include \n#include \n#include \n" + +A_LOGITS = (" const float * logits = llama_get_logits(ctx);" + " // canvas rows packed: [C or max_length, n_vocab]\n") +INS_LOGITS = A_LOGITS + """ + // DG_EB_DUMP: per-step logits + decisions (SK Stage-2 parity gate) + const char * dg_dump = getenv("DG_EB_DUMP"); + std::vector dg_canvas_in; + if (dg_dump) { + if (step_idx == 0) { + char hp[1024]; snprintf(hp, sizeof(hp), "%s/header.bin", dg_dump); + FILE * hf = fopen(hp, "wb"); + if (hf) { + int32_t hdr[4] = { n_input, C, S, params.seed }; + fwrite(hdr, 4, 4, hf); + float fl[5] = { params.t_min, params.t_max, params.entropy_bound, + (float) params.stability_threshold, params.confidence_threshold }; + fwrite(fl, 4, 5, hf); + fwrite(input_tokens, 4, n_input, hf); + fclose(hf); + } + } + if (getenv("DG_EB_DUMP_THROTTLE") && step_idx >= 2) { + char prev[1024]; snprintf(prev, sizeof(prev), "%s/step_%03d.f32", dg_dump, step_idx - 2); + while (access(prev, F_OK) == 0) { usleep(500000); } // consumer deletes + } + char lp[1024]; snprintf(lp, sizeof(lp), "%s/step_%03d.f32", dg_dump, step_idx); + FILE * lf = fopen(lp, "wb"); + if (lf) { + fwrite(logits + (size_t) logit_off * n_vocab, sizeof(float), (size_t) C * n_vocab, lf); + fclose(lf); + } + dg_canvas_in = current_canvas; + } +""" + +A_STOP = (" prev_argmax = argmax_canvas;\n" + " prev_temp_inv = temp_inv;\n") +INS_STOP = A_STOP + """ + if (dg_dump) { + char dp[1024]; snprintf(dp, sizeof(dp), "%s/step_%03d.dec", dg_dump, step_idx); + FILE * df = fopen(dp, "wb"); + if (df) { + uint8_t fin = finished ? 1 : 0; + fwrite(&step_idx, 4, 1, df); fwrite(&cur_step, 4, 1, df); fwrite(&t, 4, 1, df); + fwrite(dg_canvas_in.data(), 4, C, df); + fwrite(u.data(), 4, C, df); + fwrite(renoise.data(), 4, C, df); + fwrite(entropy.data(), 4, C, df); + fwrite(argmax_canvas.data(), 4, C, df); + fwrite(denoiser.data(), 4, C, df); + std::vector acc8(accepted.begin(), accepted.end()); + fwrite(acc8.data(), 1, C, df); + fwrite(current_canvas.data(), 4, C, df); + fwrite(&held, 4, 1, df); + fwrite(&fin, 1, 1, df); + fwrite(&entropy_sum, 4, 1, df); + fclose(df); + } + } +""" + + +def main() -> int: + p = Path(sys.argv[1]) + src = p.read_text() + if "DG_EB_DUMP" in src: + print("already instrumented") + return 0 + for anchor, ins in ((A_INCLUDE, INS_INCLUDE), (A_LOGITS, INS_LOGITS), (A_STOP, INS_STOP)): + n = src.count(anchor) + if n != 1: + print(f"anchor not unique ({n}): {anchor[:60]!r}") + return 1 + src = src.replace(anchor, ins) + p.write_text(src) + print(f"instrumented {p}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/temp/diffgemma_s2/make_embt.py b/temp/diffgemma_s2/make_embt.py new file mode 100644 index 0000000..71a7543 --- /dev/null +++ b/temp/diffgemma_s2/make_embt.py @@ -0,0 +1,43 @@ +"""make_embt.py — one-time build of the SC soft-embedding weight: token_embd +dequantized + transposed to [d_model, vocab] f16 on disk (the per-step SC +matmul then streams it like any other weight). Mirrors the reference's +dg_ensure_sc_embT (f16 transpose of the dequantized embed). + + python3 make_embt.py --gguf model.gguf --out dg_embT_f16.bin +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +from SuperKittens.models.gemma.diffusion.config import config_from_gguf # noqa: E402 +from SuperKittens.models.gemma.diffusion.gguf_io import GGUFFile # noqa: E402 + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--gguf", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--vchunk", type=int, default=8192) + args = ap.parse_args() + + gg = GGUFFile(args.gguf) + cfg = config_from_gguf(gg.meta) + d, V = cfg.d_model, cfg.vocab_size + mm = np.memmap(args.out, dtype=np.float16, mode="w+", shape=(d, V)) + for v0 in range(0, V, args.vchunk): + v1 = min(v0 + args.vchunk, V) + chunk = gg.dequant("token_embd.weight", rows=slice(v0, v1)) # [rows, d] f32 + mm[:, v0:v1] = chunk.T.astype(np.float16) + mm.flush() + del mm + print(f"wrote [{d}, {V}] f16 -> {args.out} ({d * V * 2 / 1e9:.2f} GB)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 778499c00f4ded2a480b40b54d0e8b5216286810 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 00:09:19 -0400 Subject: [PATCH 21/31] diffgemma stage2 GATE 1 GREEN: SC subgraph verified op-for-op + empirically vs instrumented reference (sc_sig oracle cos 0.9999999 / GPU 0.9999998; SC-active GPU step-1 logits argmax 256/256 vs ref; use_sc=0 gate exact-zero confirmed) --- temp/diffgemma_s2/STATUS.md | 111 +++++++++++++++ temp/diffgemma_s2/gate1_sc.py | 177 ++++++++++++++++++++++++ temp/diffgemma_s2/gate1b_gpu.sh | 55 ++++++++ temp/diffgemma_s2/gate2_real_logits.py | 142 +++++++++++++++++++ temp/diffgemma_s2/instrument_sc_dump.py | 101 ++++++++++++++ temp/diffgemma_s2/run_gate3.sh | 26 ++++ 6 files changed, 612 insertions(+) create mode 100644 temp/diffgemma_s2/STATUS.md create mode 100644 temp/diffgemma_s2/gate1_sc.py create mode 100644 temp/diffgemma_s2/gate1b_gpu.sh create mode 100644 temp/diffgemma_s2/gate2_real_logits.py create mode 100644 temp/diffgemma_s2/instrument_sc_dump.py create mode 100644 temp/diffgemma_s2/run_gate3.sh diff --git a/temp/diffgemma_s2/STATUS.md b/temp/diffgemma_s2/STATUS.md new file mode 100644 index 0000000..ae7d57b --- /dev/null +++ b/temp/diffgemma_s2/STATUS.md @@ -0,0 +1,111 @@ +# DiffusionGemma Stage 2 — sampler + self-conditioning + e2e generation + +Continues temp/diffgemma_s1/STATUS.md (forward parity). Goal: close the +EntropyBound sampler on real reference logits, verify the self-conditioning +(SC) subgraph, and produce coherent end-to-end generations on the SK stack. + +Reference: llama.cpp PR #24423 @ c84e85af (amelia `~/llamacpp-diffg`, CPU +Release, GGML_CPU_REPACK=OFF — the build cache had drifted to REPACK=ON and +was reconfigured back before any Stage-2 runs). Model: +diffusiongemma-26B-A4B-it-Q4_K_M (read-only, amelia `~/diffgemma-gguf`). +Lab: amelia `~/sk-diffg-s2b`. + +## Gate 1 — SC subgraph verification: GREEN + +### Op-for-op review (oracle graph_ref.sc_signal/embed_tokens vs PR +`dg_canvas_embed`, src/models/diffusion-gemma.cpp) + +| PR op | SK oracle | note | +|---|---|---| +| `soft_max(scale(sc_logits, sc_temp_inv))` | `softmax(sc_logits * sc_temp_inv)` | f32 | +| `mul_mat(sc_embT, probs)` (embed dequant+T, f16) | chunked `probs @ embed_rows` f32 | ggml side accumulates the 262144-long dot in f16 lanes | +| `scale(soft, sqrt(n_embd))` | `* sqrt(2816)` | | +| `build_norm(soft, sc_pre_norm, RMS)` | `rms_norm(soft, eps, self_cond_pre_norm)` | scaled rms | +| `gelu(mul_mat(sc_gate, normed))` | `gelu_tanh(normed @ gate.T)` | ggml_gelu = tanh approx | +| `mul_mat(sc_up, normed)` | `normed @ up.T` | | +| `mul_mat(sc_down, g*u)` | `(g*u) @ down.T` | | +| `scale(sc_sig, sc_use)` | `* sc_use` | {0,1} runtime gate | +| `add(canvas, sc_sig)` then `rms_norm` (no scale) | `embed_tokens(..., sc_sig)` | add BEFORE the canvas rms, both sides | + +SC temp contract (PR EB sampler): step k's forward conditions on +softmax(L_{k-1} / t_{k-1}) — `prev_temp_inv`, gated off (sc_use=0) at k=0. +sampler.py exposes exactly this (`prev_temp_inv` updated at end of step); +generate.py skips the SC compute entirely at step 0 (sig*0 == 0 bit-exactly). + +### Empirical (instrumented reference: cb(sc_sig) + cb_eval tensor dump in +the server — tools/instrument_sc_dump.py; driver tools/gate1_sc.py) + +Inputs: p1 prompt (P=20), canvas0 = sampler(seed 1234) random init, L0 = +reference zero-SC logits on canvas0, canvas1 = sampler step-0 renoise, +S=16 ⇒ t0=0.8 (sc_temp_inv=1.25). Reference forwards ~30 s each (CPU). + +| check | result | +|---|---| +| ref sc_sig at use_sc=0 | all-zero exactly (gate semantics confirmed) | +| sc_sig oracle-vs-ref [256,2816] | rel_rms 4.4e-4, cos 0.9999999, max_abs 0.014 (signal rms 5.37) | +| inp_region canvas rows oracle-vs-ref | rel_rms 4.2e-4, cos 0.99999991 | +| inp_region prompt rows oracle-vs-ref | bit-exact (0.0) | +| ref SC effect on step-1 logits (L1_sc vs L1_nosc) | rel_rms 1.35, argmax agree 99.6% — SC is a first-order input, the check has teeth | + +(GPU leg + envelope below: gate1b) + +GPU leg (forward_metal._sc_signal: probs f16 on the wire, embT [d,V] f16 +streamed in 8x352-row chunks through dg_gemm_qkt_f32, MLP on quant GEMM): + +| check | result | +|---|---| +| sc_sig GPU-vs-ref | rel_rms 6.8e-4, cos 0.99999977 | +| step-1 logits GPU-vs-ref, SC ACTIVE | argmax-canvas **256/256 = 100%**, rel_rms 0.133 | +| step-1 logits GPU-vs-ref, zero-SC control | argmax-canvas 100%, rel_rms 0.256 | +| logits finite | yes, |max| = 29.80/29.71 (softcap 30) | + +SC-active agreement is BETTER than the zero-SC envelope on the same canvas +(SC sharpens the distributions: canvas1 carries one denoise step of signal). +GPU forwards 47.5 s (zero-SC) / 52.2 s (SC) warm — the SC stream costs ~5 s. + +Benign noise note: numpy-on-Accelerate raises FP flags +(divide-by-zero/overflow/invalid "in matmul") in the host router matmul and +in f64 comparison dots; outputs verified finite + softcap-bounded, and the +SC-active forward still lands argmax-100% vs the reference. Flag noise, not +data corruption (subnormal-class inputs to BLAS kernels). + +## Gate 2 — sampler parity on REAL reference logits + +GATE2_PLACEHOLDER + +## Gate 3 — e2e coherent generation + +GATE3_PLACEHOLDER + +## Gate 4 — cross-check vs llama-diffusion-cli + +GATE4_PLACEHOLDER + +## Stage-3 baseline numbers + +GATE5_PLACEHOLDER + +## Tools (this dir, all run on amelia from ~/sk-diffg-s2b) + +- `instrument_diffusion.py` — per-step logits+decision dumps in the reference + EB sampler (DG_EB_DUMP / DG_EB_DUMP_THROTTLE keeps ≤ 2×268 MB on disk). +- `instrument_sc_dump.py` — names sc_sig in the model graph + cb_eval named- + tensor dumps in diffusion-gemma-server (DG_DUMP_TENSORS/DG_DUMP_DIR). +- `gate1_sc.py` / `gate1b_gpu.sh` — Gate-1 drivers (server + oracle + GPU). +- `gate2_real_logits.py` — spawns the instrumented cli, streams its dumps + through the SK sampler, diffs every decision field, deletes consumed logits. +- `make_embt.py` — one-time [d_model, vocab] f16 transposed embed (1.48 GB, + amelia ~/sk-diffg-s2b/dg_embT_f16.bin) for the GPU SC soft-embed stream. +- `run_gate3.sh` — watchdogged e2e generation (swap>5.5G / disk<400M / 3600s). +- `compare_eb.py`, `eb_ref_harness.cpp`, `rng_dump.cpp` — Stage-2a synthetic + sampler gate (committed earlier, still pass). + +## Host notes + +- amelia root volume runs ~2.5-4.8 GB free with the embT + dumps in place; + every logits file is 268 MB — delete as consumed (gate2 streams + deletes). +- A GGUF→derek transfer (`cat ~/diffgemma-gguf/...gguf`) was running through + amelia during the correctness runs; timing-sensitive numbers (Gate 5) were + taken TRANSFER_NOTE_PLACEHOLDER. +- Stale `/Users/amelia/SuperKittens` partial copy exists; the lab runs pin + PYTHONPATH=~/sk-diffg-s2b so the rsynced tree wins. Don't import without it. diff --git a/temp/diffgemma_s2/gate1_sc.py b/temp/diffgemma_s2/gate1_sc.py new file mode 100644 index 0000000..99adb8e --- /dev/null +++ b/temp/diffgemma_s2/gate1_sc.py @@ -0,0 +1,177 @@ +"""gate1_sc.py — Stage-2 Gate 1: SC subgraph verification on the CPU oracle +against the instrumented reference (llama-diffusion-gemma-server with +DG_DUMP_TENSORS=sc_sig,inp_region). + +Runs ON amelia. Sequence (one process at a time): + reqA: [prompt|canvas0] use_sc=0 temp=t0 -> L0 + sc_sig_r000 (must be zeros) + SK sampler(seed) consumes L0 -> canvas1, t1 + reqB: [prompt|canvas1] use_sc=1 temp=t1 -> L1_sc + sc_sig_r001 + inp_region_r001 + reqC: [prompt|canvas1] use_sc=0 temp=t1 -> L1_nosc (envelope control) +then compares the SK CPU oracle's sc_signal / embed_tokens (and optionally a +detached forward_cpu) against the dumps. + + python3 gate1_sc.py --server BIN --gguf G --prompt p1_prompt.i32 \ + --out-dir ~/sk-diffg-s2b/gate1 [--S 16 --seed 1234] [--skip-server] +""" +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from SuperKittens.models.gemma.diffusion.config import config_from_gguf # noqa: E402 +from SuperKittens.models.gemma.diffusion.gguf_io import GGUFFile # noqa: E402 +from SuperKittens.models.gemma.diffusion.graph_ref import ( # noqa: E402 + Weights, embed_tokens, sc_signal) +from SuperKittens.models.gemma.diffusion.sampler import ( # noqa: E402 + EBParams, EntropyBoundSampler) + +F32 = np.float32 + + +def stats(name: str, a: np.ndarray, b: np.ndarray) -> dict: + """a = candidate, b = reference; per-element f32 arrays, same shape.""" + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + denom = np.sqrt((b ** 2).mean()) or 1.0 + rel_rms = float(np.sqrt(((a - b) ** 2).mean()) / denom) + cos = float((a @ b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) + return {"name": name, "rel_rms": round(rel_rms, 6), "cos": round(cos, 8), + "max_abs_diff": round(float(np.abs(a - b).max()), 6)} + + +def argmax_agree(la: np.ndarray, lb: np.ndarray) -> float: + return float((la.argmax(axis=1) == lb.argmax(axis=1)).mean()) + + +class Server: + def __init__(self, binary: str, gguf: str, dump_dir: str): + env = dict(os.environ, DG_DUMP_TENSORS="sc_sig,inp_region", + DG_DUMP_DIR=dump_dir) + self.proc = subprocess.Popen([binary, gguf], stdin=subprocess.PIPE, + stdout=subprocess.PIPE, env=env, text=True) + line = self.proc.stdout.readline().strip() + assert line.startswith("READY"), line + self.n_vocab = int(line.split()[1]) + print(f"[server] {line}", flush=True) + + def forward(self, path: Path, P: int, C: int, ids: np.ndarray, + use_sc: int, temp: float) -> np.ndarray: + hdr = np.empty(4, np.int32) + hdr[0], hdr[1], hdr[2] = P, C, use_sc + hdr[3:4].view(np.float32)[0] = temp + with open(path, "wb") as f: + hdr.tofile(f) + ids.astype(np.int32).tofile(f) + t0 = time.time() + self.proc.stdin.write(str(path) + "\n") + self.proc.stdin.flush() + line = self.proc.stdout.readline().strip() + assert line == f"OK {C}", line + print(f"[server] {path.name}: {line} ({time.time() - t0:.1f}s)", flush=True) + out = np.fromfile(str(path) + ".resp", dtype=np.float32).reshape(C, -1) + os.unlink(str(path) + ".resp") # 268 MB each; the caller persists what it needs + return out + + def close(self): + try: + self.proc.stdin.write("QUIT\n") + self.proc.stdin.flush() + self.proc.wait(timeout=60) + except Exception: + self.proc.kill() + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--server", required=True) + ap.add_argument("--gguf", required=True) + ap.add_argument("--prompt", required=True) + ap.add_argument("--out-dir", required=True) + ap.add_argument("--S", type=int, default=16) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--skip-server", action="store_true", + help="reuse existing dumps/resp files in out-dir") + args = ap.parse_args() + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + prompt = np.fromfile(args.prompt, dtype=np.int32) + P = len(prompt) + + gg = GGUFFile(args.gguf) + cfg = config_from_gguf(gg.meta) + C, V = cfg.canvas_length, cfg.vocab_size + + smp = EntropyBoundSampler(EBParams(max_steps=args.S, seed=args.seed), V, C) + canvas0 = smp.canvas.copy() + t0_temp = float(smp.temperature(0)) + t1_temp = float(smp.temperature(1)) + + if not args.skip_server: + srv = Server(args.server, args.gguf, str(out)) + assert srv.n_vocab == V + ids0 = np.concatenate([prompt, canvas0]) + L0 = srv.forward(out / "reqA.bin", P, C, ids0, use_sc=0, temp=t0_temp) + L0.tofile(out / "L0.f32") + step = smp.step(np.ascontiguousarray(L0, F32)) + canvas1 = step.canvas_next + canvas1.astype(np.int32).tofile(out / "canvas1.i32") + ids1 = np.concatenate([prompt, canvas1]) + L1_sc = srv.forward(out / "reqB.bin", P, C, ids1, use_sc=1, temp=t1_temp) + L1_sc.tofile(out / "L1_sc.f32") + # reqC resets the server's sc_cache to L1_nosc, so it must come last + L1_nosc = srv.forward(out / "reqC.bin", P, C, ids1, use_sc=0, temp=t1_temp) + L1_nosc.tofile(out / "L1_nosc.f32") + srv.close() + else: + L0 = np.fromfile(out / "L0.f32", dtype=np.float32).reshape(C, V) + step = smp.step(np.ascontiguousarray(L0, F32)) + canvas1 = step.canvas_next + L1_sc = np.fromfile(out / "L1_sc.f32", dtype=np.float32).reshape(C, V) + L1_nosc = np.fromfile(out / "L1_nosc.f32", dtype=np.float32).reshape(C, V) + ids1 = np.concatenate([prompt, canvas1]) + + report: list[dict] = [] + + # 1. zero-SC request must have sc_sig == 0 exactly (the sc_use gate) + sig0 = np.fromfile(out / "sc_sig_r000.f32", dtype=np.float32) + report.append({"name": "ref sc_sig@use_sc=0 all-zero", + "ok": bool((sig0 == 0).all()), + "max_abs": float(np.abs(sig0).max())}) + + # 2. oracle sc_signal vs reference sc_sig (the SC subgraph in isolation) + w = Weights(gg) + sc_ti = F32(1.0 / t0_temp) + sig_oracle = sc_signal(w, cfg, np.ascontiguousarray(L0, F32), sc_ti, 1.0) + sig_ref = np.fromfile(out / "sc_sig_r001.f32", dtype=np.float32).reshape(C, cfg.d_model) + report.append(stats("sc_sig oracle-vs-ref", sig_oracle, sig_ref)) + report.append({"name": "sc_sig scale", "ref_rms": float(np.sqrt((sig_ref ** 2).mean())), + "oracle_rms": float(np.sqrt((sig_oracle ** 2).mean()))}) + + # 3. oracle embed (region + SC + rms) vs reference inp_region canvas rows + x_oracle = embed_tokens(w, cfg, ids1, P, sc_sig=sig_oracle) + inp_ref = np.fromfile(out / "inp_region_r001.f32", dtype=np.float32).reshape(P + C, cfg.d_model) + report.append(stats("inp_region[P:] oracle-vs-ref", x_oracle[P:], inp_ref[P:])) + report.append(stats("inp_region[:P] oracle-vs-ref", x_oracle[:P], inp_ref[:P])) + + # 4. SC-active vs zero-SC end-logits: how big is the SC effect in the ref? + report.append({"name": "ref SC effect (L1_sc vs L1_nosc)", + "argmax_agree": argmax_agree(L1_sc, L1_nosc), + "rel_rms": stats("", L1_sc, L1_nosc)["rel_rms"]}) + + for r in report: + print(json.dumps(r), flush=True) + (out / "gate1_report.json").write_text(json.dumps(report, indent=1)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/temp/diffgemma_s2/gate1b_gpu.sh b/temp/diffgemma_s2/gate1b_gpu.sh new file mode 100644 index 0000000..a84a10d --- /dev/null +++ b/temp/diffgemma_s2/gate1b_gpu.sh @@ -0,0 +1,55 @@ +#!/bin/zsh +# gate1b_gpu.sh — Gate 1 GPU leg: SK Metal forward with SC active vs the +# reference dumps produced by gate1_sc.py (same canvas1/L0 inputs). +set -e +LAB=~/sk-diffg-s2b +GGUF=~/diffgemma-gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf +G1=$LAB/gate1 +cd $LAB +export PYTHONPATH=$LAB + +echo "== GPU zero-SC control ==" +caffeinate -is python3 -m SuperKittens.models.gemma.diffusion.runner \ + --gguf $GGUF --prompt-ids ~/sk-diffg-s1/inputs/p1_prompt.i32 \ + --canvas-ids $G1/canvas1.i32 --out $G1/G1_nosc.f32 --mode gpu + +echo "== GPU SC-active (temp_inv = 1/0.8) ==" +caffeinate -is python3 -m SuperKittens.models.gemma.diffusion.runner \ + --gguf $GGUF --prompt-ids ~/sk-diffg-s1/inputs/p1_prompt.i32 \ + --canvas-ids $G1/canvas1.i32 --out $G1/G1_sc.f32 --mode gpu \ + --sc-logits $G1/L0.f32 --sc-temp-inv 1.25 \ + --sc-embt $LAB/dg_embT_f16.bin --dump-dir $G1/dump_gpu --dump-names sc_sig + +python3 - <<'EOF' +import json +import numpy as np + +G1 = "/Users/amelia/sk-diffg-s2b/gate1" +C, V, D = 256, 262144, 2816 + + +def stats(name, a, b): + a = a.astype(np.float64).ravel(); b = b.astype(np.float64).ravel() + rel = float(np.sqrt(((a - b) ** 2).mean()) / (np.sqrt((b ** 2).mean()) or 1.0)) + cos = float((a @ b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) + print(json.dumps({"name": name, "rel_rms": round(rel, 6), "cos": round(cos, 8)})) + + +def agree(name, a, b): + am = (a.argmax(1) == b.argmax(1)).mean() + print(json.dumps({"name": name, "argmax_agree": float(am)})) + + +sig_gpu = np.load(f"{G1}/dump_gpu/sc_sig.-1.npy") +sig_ref = np.fromfile(f"{G1}/sc_sig_r001.f32", np.float32).reshape(C, D) +stats("sc_sig GPU-vs-ref", sig_gpu, sig_ref) +L1_sc = np.fromfile(f"{G1}/L1_sc.f32", np.float32).reshape(C, V) +L1_nosc = np.fromfile(f"{G1}/L1_nosc.f32", np.float32).reshape(C, V) +G_sc = np.fromfile(f"{G1}/G1_sc.f32", np.float32).reshape(C, V) +G_nosc = np.fromfile(f"{G1}/G1_nosc.f32", np.float32).reshape(C, V) +agree("logits GPU-vs-ref SC-active", G_sc, L1_sc) +agree("logits GPU-vs-ref zero-SC (envelope)", G_nosc, L1_nosc) +stats("logits GPU-vs-ref SC-active", G_sc, L1_sc) +stats("logits GPU-vs-ref zero-SC (envelope)", G_nosc, L1_nosc) +EOF +echo GATE1B_DONE diff --git a/temp/diffgemma_s2/gate2_real_logits.py b/temp/diffgemma_s2/gate2_real_logits.py new file mode 100644 index 0000000..f83c765 --- /dev/null +++ b/temp/diffgemma_s2/gate2_real_logits.py @@ -0,0 +1,142 @@ +"""gate2_real_logits.py — Stage-2 Gate 2: SK sampler parity on REAL reference +logits. Spawns the instrumented llama-diffusion-cli (DG_EB_DUMP + +DG_EB_DUMP_THROTTLE, see instrument_diffusion.py) and streams its per-step +dumps: for every step, replay the SK EntropyBoundSampler on the exact logits +the reference sampler consumed and diff every decision field; delete each +268 MB logits file once consumed (disk stays <= 2 steps deep). + + python3 gate2_real_logits.py --cli BIN --gguf G --prompt "..." \ + --steps 10 --seed 1234 --dump-dir d --log out.jsonl +""" +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from SuperKittens.models.gemma.diffusion.sampler import ( # noqa: E402 + EBParams, EntropyBoundSampler) +from tools.compare_eb import read_dec # noqa: E402 + +F32 = np.float32 + + +def wait_for(path: Path, proc, timeout: float = 600.0) -> bool: + t0 = time.time() + while time.time() - t0 < timeout: + if path.exists() and path.stat().st_size > 0: + time.sleep(0.5) # writer is not atomic; settle + return True + if proc.poll() is not None: + return path.exists() + time.sleep(1.0) + return False + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--cli", required=True) + ap.add_argument("--gguf", required=True) + ap.add_argument("--prompt", required=True) + ap.add_argument("--steps", type=int, default=10) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--dump-dir", required=True) + ap.add_argument("--n-predict", type=int, default=256) + args = ap.parse_args() + + dump = Path(args.dump_dir) + dump.mkdir(parents=True, exist_ok=True) + for f in dump.glob("step_*"): + f.unlink() + (dump / "header.bin").unlink(missing_ok=True) + + env = dict(os.environ, DG_EB_DUMP=str(dump), DG_EB_DUMP_THROTTLE="1") + cli_log = open(dump / "cli.log", "w") + proc = subprocess.Popen( + [args.cli, "-m", args.gguf, "-p", args.prompt, + "--diffusion-eb-max-steps", str(args.steps), + "--seed", str(args.seed), "-n", str(args.n_predict), "-st"], + stdout=cli_log, stderr=subprocess.STDOUT, + stdin=subprocess.DEVNULL, env=env) + + if not wait_for(dump / "header.bin", proc, timeout=900): + print("FATAL: no header.bin (cli died?)", flush=True) + proc.kill() + return 2 + hdr = np.fromfile(dump / "header.bin", dtype=np.int32, count=4) + n_input, C, S, seed = (int(v) for v in hdr) + fl = np.fromfile(dump / "header.bin", dtype=np.float32, offset=16, count=5) + p = EBParams(max_steps=S, t_min=float(fl[0]), t_max=float(fl[1]), + entropy_bound=float(fl[2]), stability_threshold=int(fl[3]), + confidence_threshold=float(fl[4]), seed=seed) + print(f"header: n_input={n_input} C={C} S={S} seed={seed} " + f"t=[{fl[0]:.3f},{fl[1]:.3f}] bound={fl[2]} stab={int(fl[3])} " + f"conf={fl[4]}", flush=True) + + n_vocab = 262144 + smp = EntropyBoundSampler(p, n_vocab, C) + fields = ["canvas_in", "u", "renoise", "entropy", "argmax", "denoiser", + "accepted", "canvas_next", "held", "finished"] + total = mism = 0 + per_field: dict[str, int] = {f: 0 for f in fields} + + for s in range(S): + fdec = dump / f"step_{s:03d}.dec" + flog = dump / f"step_{s:03d}.f32" + if not wait_for(fdec, proc): + print(f"step {s}: no dec record (cli exit={proc.poll()}) — stop", + flush=True) + break + ref = read_dec(fdec, C) + canvas_before = smp.canvas.copy() + logits = np.fromfile(flog, dtype=np.float32).reshape(C, n_vocab) + r = smp.step(logits) + flog.unlink() # unblock the throttled producer + mine = {"canvas_in": canvas_before, "u": r.u, "renoise": r.renoise, + "entropy": r.entropy, "argmax": r.argmax, "denoiser": r.sampled, + "accepted": r.accepted, "canvas_next": r.canvas_next, + "held": r.held, "finished": r.finished} + line = [f"step {s:3d} t={r.t:.4f}"] + for f in fields: + a, b = mine[f], ref[f] + if f == "entropy": + d = float(np.abs(np.asarray(a) - np.asarray(b)).max()) + ok = d < 5e-4 + line.append(f"H~{d:.2e}") + elif isinstance(a, (int, bool, np.bool_)): + ok = bool(a == b) + if not ok: + line.append(f"{f}:{a}!={b}") + else: + ok = bool(np.array_equal(a, b)) + if not ok: + bad = np.nonzero(np.asarray(a) != np.asarray(b))[0] + line.append(f"{f}:{len(bad)}bad@{bad[:4].tolist()}") + total += 1 + if not ok: + mism += 1 + per_field[f] += 1 + line.append(f"acc={int(r.accepted.sum())} held={r.held} " + f"Hbar={r.entropy_mean:.4f} fin={r.finished}") + print(" ".join(line), flush=True) + if ref["finished"]: + break + + proc.wait(timeout=600) + print(f"\nfields compared: {total}, mismatched: {mism}", flush=True) + for f, n in per_field.items(): + if n: + print(f" {f}: {n} steps mismatched", flush=True) + print("TOKEN-IDENTICAL" if mism == 0 else "MISMATCH", flush=True) + return 0 if mism == 0 else 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/temp/diffgemma_s2/instrument_sc_dump.py b/temp/diffgemma_s2/instrument_sc_dump.py new file mode 100644 index 0000000..7b9afe5 --- /dev/null +++ b/temp/diffgemma_s2/instrument_sc_dump.py @@ -0,0 +1,101 @@ +"""instrument_sc_dump.py — env-driven named-tensor dumps for the Gate-1 SC +verification, patched into the predecessor's llama.cpp clone: + +1. src/models/diffusion-gemma.cpp : cb(sc_sig, "sc_sig", -1) so the SC signal + is a named graph node ("inp_region" is already named). +2. examples/diffusion-gemma-server/diffusion-gemma-server.cpp : + DG_DUMP_TENSORS= DG_DUMP_DIR= installs a sched eval + callback that writes each named tensor to /_r%03d.f32 per + request (f32, ggml layout). + +Idempotent. Usage: python3 instrument_sc_dump.py +""" +import sys +from pathlib import Path + +A_MODEL = (' sc_sig = ggml_scale(ctx0, sc_sig, dmodel.sc_use); ' + '// runtime {0,1} gate (0 == first step)\n') +INS_MODEL = A_MODEL + ' cb(sc_sig, "sc_sig", -1);\n' + +A_INC = '#include "llama.h"\n' +INS_INC = A_INC + '#include "ggml-backend.h"\n' + +A_CB = "int main(int argc, char ** argv) {\n" +INS_CB = """ +// DG_DUMP_TENSORS / DG_DUMP_DIR: per-request named-tensor dumps (SK Stage-2 SC gate) +static std::vector g_dg_dump_names; +static std::string g_dg_dump_dir; +static int g_dg_req_idx = -1; +static bool dg_dump_cb(struct ggml_tensor * t, bool ask, void * user_data) { + (void) user_data; + bool want = false; + for (const auto & s : g_dg_dump_names) { if (s == t->name) { want = true; break; } } + if (ask) { return want; } + if (want) { + const size_t n = (size_t) ggml_nelements(t); + std::vector host(n); + if (t->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(t, host.data(), 0, n * sizeof(float)); + } else { + std::vector raw(ggml_nbytes(t)); + ggml_backend_tensor_get(t, raw.data(), 0, raw.size()); + ggml_get_type_traits(t->type)->to_float(raw.data(), host.data(), (int64_t) n); + } + char p[1024]; + snprintf(p, sizeof(p), "%s/%s_r%03d.f32", g_dg_dump_dir.c_str(), t->name, g_dg_req_idx); + FILE * f = fopen(p, "wb"); + if (f) { fwrite(host.data(), sizeof(float), n, f); fclose(f); } + } + return true; +} + +""" + A_CB + +A_CTX = " llama_context * ctx = llama_init_from_model(model, cparams);\n" +INS_CTX = """ if (const char * dt = getenv("DG_DUMP_TENSORS")) { + g_dg_dump_dir = getenv("DG_DUMP_DIR") ? getenv("DG_DUMP_DIR") : "."; + std::string s = dt; + size_t pos = 0; + while (true) { + size_t e = s.find(',', pos); + g_dg_dump_names.push_back(s.substr(pos, e == std::string::npos ? std::string::npos : e - pos)); + if (e == std::string::npos) break; + pos = e + 1; + } + fprintf(stderr, "dg_dump: %zu tensors -> %s\\n", g_dg_dump_names.size(), g_dg_dump_dir.c_str()); + cparams.cb_eval = dg_dump_cb; + cparams.cb_eval_user_data = nullptr; + } +""" + A_CTX + +A_REQ = " std::vector req = read_i32_file(line);\n" +INS_REQ = A_REQ + " g_dg_req_idx++;\n" + + +def patch(path: Path, pairs) -> int: + src = path.read_text() + if "DG_DUMP_TENSORS" in src or 'cb(sc_sig, "sc_sig", -1)' in src: + print(f"already instrumented: {path}") + return 0 + for anchor, ins in pairs: + n = src.count(anchor) + if n != 1: + print(f"anchor not unique ({n}) in {path}: {anchor[:60]!r}") + return 1 + src = src.replace(anchor, ins) + path.write_text(src) + print(f"instrumented {path}") + return 0 + + +def main() -> int: + root = Path(sys.argv[1]) + rc = patch(root / "src/models/diffusion-gemma.cpp", [(A_MODEL, INS_MODEL)]) + if rc: + return rc + return patch(root / "examples/diffusion-gemma-server/diffusion-gemma-server.cpp", + [(A_INC, INS_INC), (A_CB, INS_CB), (A_CTX, INS_CTX), (A_REQ, INS_REQ)]) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/temp/diffgemma_s2/run_gate3.sh b/temp/diffgemma_s2/run_gate3.sh new file mode 100644 index 0000000..170369e --- /dev/null +++ b/temp/diffgemma_s2/run_gate3.sh @@ -0,0 +1,26 @@ +#!/bin/zsh +# run_gate3.sh [extra generate.py args...] +# One SK e2e generation, detached-safe, with the Stage-1 watchdog rules: +# kill on system swap > 5.5 GB, root-disk free < 400 MB, or wall > 3600 s. +LAB=~/sk-diffg-s2b +GGUF=~/diffgemma-gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf +name=$1; pids=$2; steps=$3; seed=$4; shift 4 +out=$LAB/gen/$name +mkdir -p $out +cd $LAB +PYTHONPATH=$LAB caffeinate -is python3 -m SuperKittens.models.gemma.diffusion.generate \ + --gguf $GGUF --prompt-ids $pids --out-dir $out --steps $steps --seed $seed \ + --sc-embt $LAB/dg_embT_f16.bin "$@" > $out/run.log 2>&1 & +PID=$! +echo $PID > $out/pid +SECS=0 +while kill -0 $PID 2>/dev/null; do + sleep 30; SECS=$((SECS+30)) + swap=$(sysctl -n vm.swapusage | awk '{print $6}' | tr -d M) + free=$(df -m /System/Volumes/Data | tail -1 | awk '{print $4}') + if (( ${swap%.*} > 5500 )); then echo "WATCHDOG swap=${swap}M KILL" >> $out/run.log; kill -9 $PID; exit 3; fi + if (( free < 400 )); then echo "WATCHDOG diskfree=${free}M KILL" >> $out/run.log; kill -9 $PID; exit 4; fi + if (( SECS > 3600 )); then echo "WATCHDOG timeout KILL" >> $out/run.log; kill -9 $PID; exit 5; fi +done +wait $PID +echo "EXIT $?" >> $out/run.log From 0bd32342eb7f64b08c3001e8adc947d28986d6ed Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 00:15:58 -0400 Subject: [PATCH 22/31] diffgemma stage2 GATE 2 GREEN: sampler TOKEN-IDENTICAL on real reference logits (instrumented cli, S=10 seed=1234 C=256; 7-step trajectory incl. adaptive stop, 0/70 decision fields differ, max|dH| 2.4e-7) --- temp/diffgemma_s2/STATUS.md | 29 +++++++++++++++++++++++++- temp/diffgemma_s2/gate2_real_logits.py | 2 +- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/temp/diffgemma_s2/STATUS.md b/temp/diffgemma_s2/STATUS.md index ae7d57b..41cb09a 100644 --- a/temp/diffgemma_s2/STATUS.md +++ b/temp/diffgemma_s2/STATUS.md @@ -71,7 +71,34 @@ data corruption (subnormal-class inputs to BLAS kernels). ## Gate 2 — sampler parity on REAL reference logits -GATE2_PLACEHOLDER +**GREEN — TOKEN-IDENTICAL, 0/70 fields mismatched.** + +Setup: instrumented `llama-diffusion-cli` (DG_EB_DUMP, throttle ≤ 2 logits +files on disk), prompt "What is the capital of France?" through the cli's own +chat template (P=23: `<|turn>system\n<|think|>\n\n<|turn>user\n…`), +S=10, seed=1234, C=256, CPU reference w/ prompt-KV cache. Consumer +(tools/gate2_real_logits.py) replays the SK sampler on the exact per-step +logits the reference consumed and diffs EVERY decision field. + +| step | t | max\|ΔH\| | accepted | held | Hbar | fin | +|---|---|---|---|---|---|---| +| 0 | 0.80 | 2.4e-7 | 122 | 0 | 0.468 | no | +| 1 | 0.76 | 2.4e-7 | 189 | 0 | 0.311 | no | +| 2 | 0.72 | 2.4e-7 | 205 | 0 | 0.265 | no | +| 3 | 0.68 | 1.2e-7 | 214 | 0 | 0.123 | no | +| 4 | 0.64 | 1.2e-7 | 230 | 0 | 0.068 | no | +| 5 | 0.60 | 6.0e-8 | 243 | 0 | 0.0085 | no | +| 6 | 0.56 | 1.5e-8 | 255 | 1 | 0.0016 | **yes** | + +All of: working canvas, RNG draws (u, renoise), entropy (≤2.4e-7), argmax, +multinomial picks, accept-sets, renoised canvas, held counter, and the +adaptive stop — identical on real logits, every step. The documented +cum-plateau flip risk did not materialize on this trajectory (it remains +possible in principle on exact f32 ties; the synthetic gate quantified it at +2/1792 discarded picks). + +Reference cli's own final answer (same run): thought channel reasoning + +"The capital of France is Paris." — 7 steps, 30.5 s/step CPU. ## Gate 3 — e2e coherent generation diff --git a/temp/diffgemma_s2/gate2_real_logits.py b/temp/diffgemma_s2/gate2_real_logits.py index f83c765..9c523cc 100644 --- a/temp/diffgemma_s2/gate2_real_logits.py +++ b/temp/diffgemma_s2/gate2_real_logits.py @@ -62,7 +62,7 @@ def main() -> int: proc = subprocess.Popen( [args.cli, "-m", args.gguf, "-p", args.prompt, "--diffusion-eb-max-steps", str(args.steps), - "--seed", str(args.seed), "-n", str(args.n_predict), "-st"], + "--seed", str(args.seed), "-n", str(args.n_predict)], stdout=cli_log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, env=env) From c6b94fd6c94c9567157756522594fbaddedb64df Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 00:24:58 -0400 Subject: [PATCH 23/31] granite: hybrid family package (mamba2+attention interleave, NoPE, granite multipliers) granite-4.0-h-1b: 36 mamba2 + 4 attention layers (per-layer type from GGUF head_count_kv), dense SwiGLU FFN every layer, tied Q8_0 head. Reuses mamba2 family kernels (conv1d_silu/_step, mamba2_ssd, gate_norm) and shared dense kernels (q8_0_matvec, mha_causal D=128, kv_cache_write, rmsnorm, silu_mul). New granite_ops.metal carries the granite scalar multipliers; Q is pre-scaled by attention_scale*sqrt(D) so mha_causal's hardcoded 1/sqrt(D) becomes the granite attention_multiplier. h-1b over h-micro because h-micro is head_dim=64 (SK dense attention is D=128-only). --- SuperKittens/inference/registry.py | 22 + SuperKittens/models/granite/__init__.py | 3 + SuperKittens/models/granite/granite.py | 220 +++++++ SuperKittens/models/granite/granite_model.h | 539 ++++++++++++++++ SuperKittens/models/granite/granite_ops.metal | 39 ++ SuperKittens/models/granite/launcher.c++ | 574 ++++++++++++++++++ SuperKittens/models/granite/launcher.h | 78 +++ .../models/load/tokenizer/tokenizer.py | 5 + 8 files changed, 1480 insertions(+) create mode 100644 SuperKittens/models/granite/__init__.py create mode 100644 SuperKittens/models/granite/granite.py create mode 100644 SuperKittens/models/granite/granite_model.h create mode 100644 SuperKittens/models/granite/granite_ops.metal create mode 100644 SuperKittens/models/granite/launcher.c++ create mode 100644 SuperKittens/models/granite/launcher.h diff --git a/SuperKittens/inference/registry.py b/SuperKittens/inference/registry.py index 28481db..58290b9 100644 --- a/SuperKittens/inference/registry.py +++ b/SuperKittens/inference/registry.py @@ -402,6 +402,28 @@ class ModelSpec: head_dim=128, n_int=25600, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, tie_word_embeddings=0), ), + # Granite-4.0-H-1B: IBM granitehybrid — 40 layers, 36 mamba2 + 4 attention + # (layers 5/15/25/35; per-layer type from GGUF head_count_kv), dense SwiGLU + # FFN on EVERY layer. Attention is NoPE (no positional encoding) with + # attention_multiplier 1/128 replacing 1/sqrt(head_dim); embeddings/residuals/ + # logits carry granite scalar multipliers (12 / 0.22 / 1/6). head_dim=128 is + # why h-1b and not h-micro (h-micro is head_dim=64; SK dense attention is + # D=128). Tied Q8_0 head. Reuses mamba2 family kernels + shared dense kernels. + "granite-4.0-h-1b": ModelSpec( + family="granite", + adapter="SuperKittens.models.granite.granite:Granite", + hf_repo="ibm-granite/granite-4.0-h-1b", + weight_dir="granite-4.0-h-1b-GGUF", + gguf_name="granite-4.0-h-1b-Q8_0.gguf", + default_quant="q8_0", + tokenizer_family="granite", + dims=dict(n_layers=40, d_model=1536, n_heads=12, n_kv_heads=4, + head_dim=128, n_int=4096, d_inner=3072, ssm_n_heads=48, + ssm_head_dim=64, ssm_state=128, ssm_n_groups=1, ssm_conv=4, + vocab_size=100352, eps=1e-5, + embedding_scale=12.0, residual_scale=0.22, + attention_scale=0.0078125, logit_scale=6.0), + ), # DiffusionGemma 26B-A4B: block text-diffusion MoE on a gemma4 backbone # (llama.cpp PR #24423 is the runtime reference). Stage-1 adapter exposes # the unified zero-SC forward only; the entropy-bound sampler is Stage 2. diff --git a/SuperKittens/models/granite/__init__.py b/SuperKittens/models/granite/__init__.py new file mode 100644 index 0000000..747f7d1 --- /dev/null +++ b/SuperKittens/models/granite/__init__.py @@ -0,0 +1,3 @@ +from .granite import Granite, Config + +__all__ = ["Granite", "Config"] diff --git a/SuperKittens/models/granite/granite.py b/SuperKittens/models/granite/granite.py new file mode 100644 index 0000000..35977fa --- /dev/null +++ b/SuperKittens/models/granite/granite.py @@ -0,0 +1,220 @@ +"""granite.py — IBM Granite-4.x hybrid (interleaved mamba2 + attention) adapter. + +Drives the sk_granite_* C-ABI launcher: a hybrid stack where per-layer type +comes from GGUF metadata (head_count_kv 0 = mamba2, >0 = attention), every +layer carries a dense SwiGLU FFN, attention is NoPE with the granite +attention_multiplier, and embeddings/residuals/logits carry granite's scalar +multipliers. Reuses the mamba2 family kernels and the shared dense kernels. +""" +from __future__ import annotations +import ctypes +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from SuperKittens.inference.generation import Model +from SuperKittens.inference.c_binder import bind, optional, CtypesConfig + + +class _Config(ctypes.Structure): + _fields_ = [ + ("batch", ctypes.c_uint32), + ("seq_max", ctypes.c_uint32), + ("cache_max", ctypes.c_uint32), + ("n_layers", ctypes.c_uint32), + ("d_model", ctypes.c_uint32), + ("n_heads", ctypes.c_uint32), + ("n_kv_heads", ctypes.c_uint32), + ("head_dim", ctypes.c_uint32), + ("n_int", ctypes.c_uint32), + ("d_inner", ctypes.c_uint32), + ("ssm_n_heads", ctypes.c_uint32), + ("ssm_head_dim", ctypes.c_uint32), + ("ssm_state", ctypes.c_uint32), + ("ssm_n_groups", ctypes.c_uint32), + ("ssm_conv", ctypes.c_uint32), + ("vocab_size", ctypes.c_uint32), + ("eps", ctypes.c_float), + ("embedding_scale", ctypes.c_float), + ("residual_scale", ctypes.c_float), + ("attention_scale", ctypes.c_float), + ("logit_scale", ctypes.c_float), + ] + + +GRANITE_ABI = { + "create": ([ctypes.POINTER(_Config)], ctypes.c_void_p), + "load_gguf": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "forward": ([ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32)], ctypes.c_int), + "generate_n": ([ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.c_int32], ctypes.c_int), + "get_last_logits": ([ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), + "get_pos": optional([ctypes.c_void_p], ctypes.c_uint32), + "reset": ([ctypes.c_void_p], None), + "destroy": ([ctypes.c_void_p], None), +} + +_lib = None +def _load(): + global _lib + if _lib is None: + _lib = bind("granite", GRANITE_ABI) + return _lib + + +@dataclass +class Config(CtypesConfig): + # granite-4.0-h-1b defaults (verified against the GGUF metadata). + n_layers: int = 40 + d_model: int = 1536 + n_heads: int = 12 + n_kv_heads: int = 4 + head_dim: int = 128 + n_int: int = 4096 + d_inner: int = 3072 + ssm_n_heads: int = 48 + ssm_head_dim: int = 64 + ssm_state: int = 128 + ssm_n_groups: int = 1 + ssm_conv: int = 4 + vocab_size: int = 100352 + eps: float = 1e-5 + embedding_scale: float = 12.0 + residual_scale: float = 0.22 + attention_scale: float = 0.0078125 + logit_scale: float = 6.0 + batch: int = 1 + seq_max: int = 1024 + cache_max: int = 4096 + + +class Granite(Model): + _repr_fields = (("L", "n_layers"), ("D", "d_model"), ("E", "d_inner"), + ("ssmH", "ssm_n_heads"), ("attnH", "n_heads")) + + def __init__(self, config: Config | None = None): + self.cfg = config or Config() + lib = _load() + self._destroy_fn = lib.sk_granite_destroy + self._cstruct = self.cfg.to_c(_Config) + self._h = lib.sk_granite_create(ctypes.byref(self._cstruct)) + if not self._h: + raise RuntimeError("sk_granite_create failed (missing PSO?)") + self._last_token = None + self.tokenizer = None + self.vocab_size = self.cfg.vocab_size + + def load_gguf(self, path: str) -> None: + rc = _load().sk_granite_load_gguf(self._h, str(path).encode()) + if rc: + raise RuntimeError(f"sk_granite_load_gguf failed: {rc}") + + def reset(self) -> None: + _load().sk_granite_reset(self._h) + self._last_token = None + + def _forward(self, input_ids: np.ndarray) -> np.ndarray: + ids = np.ascontiguousarray(np.asarray(input_ids, dtype=np.int32)).reshape(-1) + out = np.empty((1,), dtype=np.int32) + rc = _load().sk_granite_forward( + self._h, + ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)), + ids.size, + out.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))) + if rc: + raise RuntimeError(f"sk_granite_forward failed: {rc}") + self._last_token = int(out[0]) + return out + + def forward(self, input_ids) -> int: + return int(self._forward(input_ids)[0]) + + def _last_logits(self) -> np.ndarray: + out = np.empty((self.cfg.vocab_size,), dtype=np.float16) + rc = _load().sk_granite_get_last_logits(self._h, out.ctypes.data) + if rc: + raise RuntimeError(f"sk_granite_get_last_logits failed: {rc}") + return out + + def generate(self, input_ids, *, max_new_tokens: int = 64, + temperature: float = 0.0, top_p: float = 1.0, + top_k=None, eos_id=None, eos_ids=None, + seed: int = 0, sampler=None): + greedy = (sampler is None and temperature <= 0.0 + and (top_p >= 1.0 or top_p <= 0.0) and not top_k) + if not greedy: + return super().generate(input_ids, max_new_tokens=max_new_tokens, + temperature=temperature, top_p=top_p, + top_k=top_k, eos_id=eos_id, eos_ids=eos_ids, + seed=seed, sampler=sampler) + stops = set() + if eos_ids: + stops |= {int(x) for x in eos_ids} + if eos_id is not None: + stops.add(int(eos_id)) + if not stops and self.tokenizer is not None: + t_eos = getattr(self.tokenizer, "eos_ids", None) + if t_eos: + stops |= {int(x) for x in t_eos} + eos_single = int(next(iter(stops))) if len(stops) == 1 else -1 + + ids = np.ascontiguousarray(np.asarray(input_ids, dtype=np.int32)).reshape(-1) + out = np.empty((int(max_new_tokens),), dtype=np.int32) + self.reset() + n = _load().sk_granite_generate_n( + self._h, + ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)), + ctypes.c_uint32(ids.size), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)), + ctypes.c_uint32(int(max_new_tokens)), + ctypes.c_int32(eos_single)) + if n < 0: + raise RuntimeError(f"sk_granite_generate_n failed: {n}") + toks = out[:n].tolist() + if len(stops) > 1: + for i, t in enumerate(toks): + if t in stops: + toks = toks[:i + 1] + break + self._last_token = toks[-1] if toks else None + return toks + + @classmethod + def from_spec(cls, spec, **overrides) -> "Granite": + sk_root = Path(__file__).resolve().parents[3] + snap = Path(overrides.pop("snapshot", None) + or (sk_root / "SuperKittens" / "model_weights" / spec.weight_dir)) + + from dataclasses import fields + allowed = {f.name for f in fields(Config)} + d = {k: v for k, v in dict(spec.dims).items() if k in allowed} + cfg = Config(**d) + for k, v in overrides.items(): + if hasattr(cfg, k): + setattr(cfg, k, v) + + gguf_path = snap / spec.gguf_name if spec.gguf_name else None + if not gguf_path or not gguf_path.exists(): + raise FileNotFoundError(f"granite GGUF not found: {gguf_path}") + + m = cls(cfg) + m.load_gguf(str(gguf_path)) + m._attach_tokenizer(spec, snap) + return m + + def _attach_tokenizer(self, spec, snap: Path) -> None: + if not spec.tokenizer_family: + return + try: + from SuperKittens.models.load.tokenizer import Tokenizer + json_path = snap / "tokenizer.json" + if json_path.exists(): + self.tokenizer = Tokenizer.from_hf_json(str(json_path), + family=spec.tokenizer_family) + else: + print(f"[granite] no tokenizer.json in {snap}") + except Exception as e: + print(f"[granite] tokenizer attach failed: {e}") diff --git a/SuperKittens/models/granite/granite_model.h b/SuperKittens/models/granite/granite_model.h new file mode 100644 index 0000000..ae784ce --- /dev/null +++ b/SuperKittens/models/granite/granite_model.h @@ -0,0 +1,539 @@ +// Granite-4.x hybrid dispatch orchestrator. +// +// Per layer (llama.cpp granite-hybrid.cpp is the reference contract): +// 1. RMSNorm(x, attn_norm) -> x_norm +// 2. mixer: +// mamba layer — in_proj -> [z|x|B|C|dt] split -> conv1d+SiLU (O(1) +// decode state carry) -> SSD scan -> SiLU(z)-gated RMSNorm -> out_proj +// (reuses the mamba2 family kernels verbatim) +// attn layer — q/k/v Q8_0 matvecs, NoPE (no RoPE dispatch at all), +// Q pre-scaled so mha_causal's 1/sqrt(D) becomes attention_scale, +// kv-cache write, mha_causal (D=128), o-proj +// 3. ffn_inp = x + residual_scale * mixer_out +// 4. RMSNorm(ffn_inp, ffn_norm) -> SwiGLU FFN (gate/up/silu_mul/down) +// 5. x_out = ffn_inp + residual_scale * ffn_out +// Model: embed * embedding_scale -> layers -> final RMSNorm -> tied Q8_0 head +// (last row only) -> logits / logit_scale -> argmax. + +#ifndef SK_GRANITE_MODEL_H +#define SK_GRANITE_MODEL_H + +#include +#include +#include +#include + +namespace meow { +namespace granite { + +struct PSOs { + MTL::ComputePipelineState* rmsnorm = nullptr; + MTL::ComputePipelineState* rmsnorm_t1 = nullptr; // optional T=1 fast path + MTL::ComputePipelineState* q8_0_matvec = nullptr; + MTL::ComputePipelineState* split_packed = nullptr; + MTL::ComputePipelineState* conv1d_silu = nullptr; + MTL::ComputePipelineState* conv1d_silu_step = nullptr; + MTL::ComputePipelineState* conv_state_capture = nullptr; + MTL::ComputePipelineState* mamba2_ssd = nullptr; // handles L=1 decode (state carry) + MTL::ComputePipelineState* gate_norm = nullptr; + MTL::ComputePipelineState* silu_mul = nullptr; + MTL::ComputePipelineState* attn = nullptr; // mha_causal (D=128) + MTL::ComputePipelineState* kv_cache_write = nullptr; + MTL::ComputePipelineState* t_seq_to_head = nullptr; + MTL::ComputePipelineState* t_head_to_seq = nullptr; + MTL::ComputePipelineState* scale = nullptr; // granite_scale_f16 + MTL::ComputePipelineState* add_scaled = nullptr; // granite_add_scaled_f16 + MTL::ComputePipelineState* embedding_lookup = nullptr; + MTL::ComputePipelineState* argmax = nullptr; +}; + +// Per-layer weights. Mamba and attention members are mutually exclusive +// (nullptr for the other type); FFN members exist on every layer. +struct LayerWeights { + bool is_attn = false; + // mamba2 (Q8_0 projections, fp16 small tensors) + MTL::Buffer* ssm_in = nullptr; // Q8_0 (2E+2GN+H, D) + MTL::Buffer* ssm_out = nullptr; // Q8_0 (D, E) + MTL::Buffer* conv_w = nullptr; // fp16 (C_in, K) — conv1d_silu reads w[c*K+k] + MTL::Buffer* conv_b = nullptr; // fp16 (C_in,) + MTL::Buffer* dt_bias = nullptr; // fp16 (H,) + MTL::Buffer* A_log = nullptr; // fp16 (H,) = log(-ssm_a_gguf) + MTL::Buffer* ssm_D = nullptr; // fp16 (H,) + MTL::Buffer* ssm_norm = nullptr; // fp16 (E,) + // attention (Q8_0) + MTL::Buffer* wq = nullptr; // (qN, D) + MTL::Buffer* wk = nullptr; // (kvN, D) + MTL::Buffer* wv = nullptr; // (kvN, D) + MTL::Buffer* wo = nullptr; // (D, qN) + // FFN (Q8_0, every layer) + MTL::Buffer* gate = nullptr; // (n_int, D) + MTL::Buffer* up = nullptr; // (n_int, D) + MTL::Buffer* down = nullptr; // (D, n_int) +}; + +struct LayerState { + // attention layers + MTL::Buffer* k_cache = nullptr; // (cache_max, n_kv, hd) fp16 + MTL::Buffer* v_cache = nullptr; + // mamba layers + MTL::Buffer* conv_state = nullptr; // (K-1, C_in) fp16 + MTL::Buffer* ssm_state = nullptr; // (H, P, N) fp32 +}; + +struct Weights { + MTL::Buffer* embed = nullptr; // fp16 (V, D) dequant — lookup table + MTL::Buffer* head_q8 = nullptr; // Q8_0 (V, D) raw — tied LM head matvec + MTL::Buffer* attn_norm = nullptr; // fp16 (n_layers, D) pre-mixer norm + MTL::Buffer* ffn_norm = nullptr; // fp16 (n_layers, D) + MTL::Buffer* final_norm = nullptr; // fp16 (D,) + std::vector layers; +}; + +struct Buffers { + MTL::Buffer* input_ids = nullptr; + MTL::Buffer* output_id = nullptr; + MTL::Buffer* x_a = nullptr; // residual ping + MTL::Buffer* x_b = nullptr; // residual pong + MTL::Buffer* x_norm = nullptr; + MTL::Buffer* ffn_inp = nullptr; // post-mixer residual + MTL::Buffer* mixer_out = nullptr; // (T, D) mixer block output + MTL::Buffer* logits = nullptr; // (seq_max, V) fp16 + // mamba scratch + MTL::Buffer* in_proj_out = nullptr; // (T, 2E+2GN+H) + MTL::Buffer* z = nullptr; // (T, E) + MTL::Buffer* xBC = nullptr; // (T, C_in) + MTL::Buffer* dt_raw = nullptr; // (T, H) + MTL::Buffer* xBC_post = nullptr; // (T, C_in) (also split scratch) + MTL::Buffer* ssd_out = nullptr; // (T, E) + MTL::Buffer* gated = nullptr; // (T, E) + // attention scratch + MTL::Buffer* q = nullptr; // (T, qN) + MTL::Buffer* k_tmp = nullptr; // (T, kvN) + MTL::Buffer* v_tmp = nullptr; // (T, kvN) + MTL::Buffer* attn_out = nullptr; // (T, qN) + MTL::Buffer* q_th = nullptr; // head-major prefill scratch + MTL::Buffer* k_th = nullptr; + MTL::Buffer* v_th = nullptr; + MTL::Buffer* attn_out_seq = nullptr; + // FFN scratch + MTL::Buffer* gate_buf = nullptr; // (T, n_int) + MTL::Buffer* up_buf = nullptr; // (T, n_int) + MTL::Buffer* mlp_out = nullptr; // (T, D) +}; + +struct Params { + uint32_t seq = 1; + uint32_t n_layers = 40; + uint32_t d_model = 1536; + uint32_t n_heads = 12; + uint32_t n_kv_heads = 4; + uint32_t head_dim = 128; + uint32_t n_int = 4096; + uint32_t d_inner = 3072; // E + uint32_t ssm_heads = 48; // H + uint32_t ssm_pdim = 64; // P + uint32_t ssm_state = 128; // N + uint32_t ssm_groups = 1; // G + uint32_t ssm_conv = 4; // K + uint32_t vocab_size = 100352; + uint32_t cache_max = 4096; + uint32_t current_pos = 0; + float eps = 1e-5f; + float embedding_scale = 12.0f; + float residual_scale = 0.22f; + float attention_scale = 0.0078125f; + float logit_scale = 6.0f; +}; + +// ── encode helpers (one encoder per op; encoder boundaries order producers) ── + +inline void enc_rmsnorm(MTL::CommandBuffer* cmd, const PSOs& P, + MTL::Buffer* x, MTL::Buffer* gamma, size_t off_g, + MTL::Buffer* out, uint32_t rows, uint32_t n, float eps) +{ + const bool t1 = (P.rmsnorm_t1 != nullptr) && (rows == 1u); + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(t1 ? P.rmsnorm_t1 : P.rmsnorm); + enc->setBuffer(x, 0, 0); + enc->setBuffer(gamma, off_g, 1); + enc->setBuffer(out, 0, 2); + enc->setBytes(&rows, 4, 3); + enc->setBytes(&n, 4, 4); + enc->setBytes(&eps, 4, 5); + if (t1) enc->dispatchThreadgroups(MTL::Size(1, rows, 1), MTL::Size(256, 1, 1)); + else enc->dispatchThreadgroups(MTL::Size(1, (rows + 3) / 4, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); +} + +// y[T,N] (row stride ldC) = x[T,K] @ W(Q8_0 [N,K] row-major)^T, per-row matvec. +inline void enc_q8_matvec(MTL::CommandBuffer* cmd, const PSOs& P, + MTL::Buffer* W, MTL::Buffer* x, MTL::Buffer* y, + uint32_t T, uint32_t N, uint32_t K, + size_t off_y = 0, uint32_t ldC = 0) +{ + if (ldC == 0) ldC = N; + auto* enc = cmd->computeCommandEncoder(); + for (uint32_t m = 0; m < T; ++m) { + if (m) enc->memoryBarrier(MTL::BarrierScopeBuffers); + enc->setComputePipelineState(P.q8_0_matvec); + enc->setBuffer(x, (size_t)m * K * 2, 0); + enc->setBuffer(W, 0, 1); + enc->setBuffer(y, off_y + (size_t)m * ldC * 2, 2); + enc->setBytes(&K, 4, 3); + enc->setBytes(&N, 4, 4); + enc->dispatchThreadgroups(MTL::Size((N + 1) / 2, 1, 1), MTL::Size(128, 1, 1)); + } + enc->endEncoding(); +} + +inline void enc_split(MTL::CommandBuffer* cmd, const PSOs& P, + MTL::Buffer* src, MTL::Buffer* outA, MTL::Buffer* outB, + uint32_t T, uint32_t A, uint32_t B) +{ + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.split_packed); + enc->setBuffer(src, 0, 0); + enc->setBuffer(outA, 0, 1); + enc->setBuffer(outB, 0, 2); + enc->setBytes(&T, 4, 3); + enc->setBytes(&A, 4, 4); + enc->setBytes(&B, 4, 5); + enc->dispatchThreads(MTL::Size(A + B, T, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); +} + +inline void enc_scale(MTL::CommandBuffer* cmd, const PSOs& P, + MTL::Buffer* x, float s, uint32_t n, size_t off_x = 0) +{ + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.scale); + enc->setBuffer(x, off_x, 0); + enc->setBytes(&s, 4, 1); + enc->setBytes(&n, 4, 2); + enc->dispatchThreads(MTL::Size(n, 1, 1), MTL::Size(256, 1, 1)); + enc->endEncoding(); +} + +inline void enc_add_scaled(MTL::CommandBuffer* cmd, const PSOs& P, + MTL::Buffer* a, MTL::Buffer* b, MTL::Buffer* y, + float s, uint32_t n) +{ + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.add_scaled); + enc->setBuffer(a, 0, 0); + enc->setBuffer(b, 0, 1); + enc->setBuffer(y, 0, 2); + enc->setBytes(&s, 4, 3); + enc->setBytes(&n, 4, 4); + enc->dispatchThreads(MTL::Size(n, 1, 1), MTL::Size(256, 1, 1)); + enc->endEncoding(); +} + +inline void enc_transpose(MTL::CommandBuffer* cmd, MTL::ComputePipelineState* pso, + MTL::Buffer* src, MTL::Buffer* dst, + uint32_t T, uint32_t H, uint32_t D) +{ + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(pso); + enc->setBuffer(src, 0, 0); + enc->setBuffer(dst, 0, 1); + enc->setBytes(&T, 4, 2); + enc->setBytes(&H, 4, 3); + enc->setBytes(&D, 4, 4); + enc->dispatchThreads(MTL::Size(D, T, H), MTL::Size(32, 1, 1)); + enc->endEncoding(); +} + +// ── mamba2 mixer (reuses the models/ssm/mamba2 kernels) ───────────────────── + +inline void dispatch_mamba_mixer(MTL::CommandBuffer* cmd, const PSOs& P, + const Params& p, const LayerWeights& W, + const LayerState& S, Buffers& B) +{ + const uint32_t T = p.seq; + const uint32_t D = p.d_model; + const uint32_t E = p.d_inner; + const uint32_t H = p.ssm_heads; + const uint32_t Pd = p.ssm_pdim; + const uint32_t G = p.ssm_groups; + const uint32_t N = p.ssm_state; + const uint32_t K = p.ssm_conv; + const uint32_t C_in = E + 2 * G * N; + const uint32_t IN_OUT = 2 * E + 2 * G * N + H; + const bool decode = (T == 1) && (p.current_pos > 0); + + enc_q8_matvec(cmd, P, W.ssm_in, B.x_norm, B.in_proj_out, T, IN_OUT, D); + + // [z(E) | xBC(C_in) | dt(H)] — two splits, xBC_post doubles as scratch. + enc_split(cmd, P, B.in_proj_out, B.z, B.xBC_post, T, E, C_in + H); + enc_split(cmd, P, B.xBC_post, B.xBC, B.dt_raw, T, C_in, H); + + if (decode) { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.conv1d_silu_step); + enc->setBuffer(B.xBC, 0, 0); + enc->setBuffer(W.conv_w, 0, 1); + enc->setBuffer(W.conv_b, 0, 2); + enc->setBuffer(B.xBC_post, 0, 3); + enc->setBuffer(S.conv_state, 0, 4); + const uint32_t one = 1; + enc->setBytes(&one, 4, 5); + enc->setBytes(&C_in, 4, 6); + enc->setBytes(&K, 4, 7); + enc->dispatchThreads(MTL::Size(C_in, 1, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); + } else { + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.conv1d_silu); + enc->setBuffer(B.xBC, 0, 0); + enc->setBuffer(W.conv_w, 0, 1); + enc->setBuffer(W.conv_b, 0, 2); + enc->setBuffer(B.xBC_post, 0, 3); + const uint32_t one = 1; + enc->setBytes(&one, 4, 4); + enc->setBytes(&T, 4, 5); + enc->setBytes(&C_in, 4, 6); + enc->dispatchThreadgroups(MTL::Size(1, (T + 3) / 4, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.conv_state_capture); + enc->setBuffer(B.xBC, 0, 0); + enc->setBuffer(S.conv_state, 0, 1); + const uint32_t one = 1; + enc->setBytes(&one, 4, 2); + enc->setBytes(&T, 4, 3); + enc->setBytes(&C_in, 4, 4); + enc->setBytes(&K, 4, 5); + enc->dispatchThreads(MTL::Size(C_in, K - 1, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + } + + // SSD scan: x/B/C alias xBC_post (token stride C_in). Granite has no + // time_step_limit -> (0, +inf) matches ggml's softplus-only dt. + { + const size_t off_B_in = (size_t)E * 2; + const size_t off_C_in = (size_t)(E + G * N) * 2; + const float dt_min = 0.0f, dt_max = INFINITY; + const uint32_t one = 1; + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.mamba2_ssd); + enc->setBuffer(B.xBC_post, 0, 0); + enc->setBuffer(B.dt_raw, 0, 1); + enc->setBuffer(W.A_log, 0, 2); + enc->setBuffer(B.xBC_post, off_B_in, 3); + enc->setBuffer(B.xBC_post, off_C_in, 4); + enc->setBuffer(W.ssm_D, 0, 5); + enc->setBuffer(W.dt_bias, 0, 6); + enc->setBuffer(B.ssd_out, 0, 7); + enc->setBuffer(S.ssm_state, 0, 8); + enc->setBytes(&one, 4, 9); + enc->setBytes(&T, 4, 10); + enc->setBytes(&H, 4, 11); + enc->setBytes(&Pd, 4, 12); + enc->setBytes(&G, 4, 13); + enc->setBytes(&N, 4, 14); + enc->setBytes(&dt_min, 4, 15); + enc->setBytes(&dt_max, 4, 16); + enc->setBytes(&C_in, 4, 17); + enc->dispatchThreadgroups(MTL::Size(H, Pd, 1), MTL::Size(N, 1, 1)); + enc->endEncoding(); + } + + // y = RMSNorm(ssd_out * SiLU(z)) * ssm_norm (G=1: norm over full E) + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.gate_norm); + enc->setBuffer(B.ssd_out, 0, 0); + enc->setBuffer(B.z, 0, 1); + enc->setBuffer(W.ssm_norm, 0, 2); + enc->setBuffer(B.gated, 0, 3); + enc->setBytes(&T, 4, 4); + enc->setBytes(&E, 4, 5); + enc->setBytes(&p.eps, 4, 6); + enc->dispatchThreadgroups(MTL::Size(1, (T + 3) / 4, 1), MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + + enc_q8_matvec(cmd, P, W.ssm_out, B.gated, B.mixer_out, T, D, E); +} + +// ── attention mixer (NoPE; Q pre-scaled for granite's attention_scale) ────── + +inline void dispatch_attn_mixer(MTL::CommandBuffer* cmd, const PSOs& P, + const Params& p, const LayerWeights& W, + const LayerState& S, Buffers& B) +{ + const uint32_t T = p.seq; + const uint32_t D = p.d_model; + const uint32_t hd = p.head_dim; + const uint32_t qN = p.n_heads * hd; + const uint32_t kvN = p.n_kv_heads * hd; + + enc_q8_matvec(cmd, P, W.wq, B.x_norm, B.q, T, qN, D); + enc_q8_matvec(cmd, P, W.wk, B.x_norm, B.k_tmp, T, kvN, D); + enc_q8_matvec(cmd, P, W.wv, B.x_norm, B.v_tmp, T, kvN, D); + + // mha_causal hardcodes softmax scale 1/sqrt(D); pre-scaling Q by + // attention_scale*sqrt(D) makes the effective scale attention_scale. + enc_scale(cmd, P, B.q, p.attention_scale * std::sqrt((float)hd), T * qN); + + MTL::Buffer* q_in = B.q; + MTL::Buffer* k_in = B.k_tmp; + MTL::Buffer* v_in = B.v_tmp; + if (T > 1) { + enc_transpose(cmd, P.t_seq_to_head, B.q, B.q_th, T, p.n_heads, hd); + enc_transpose(cmd, P.t_seq_to_head, B.k_tmp, B.k_th, T, p.n_kv_heads, hd); + enc_transpose(cmd, P.t_seq_to_head, B.v_tmp, B.v_th, T, p.n_kv_heads, hd); + q_in = B.q_th; k_in = B.k_th; v_in = B.v_th; + } + + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.kv_cache_write); + enc->setBuffer(k_in, 0, 0); + enc->setBuffer(v_in, 0, 1); + enc->setBuffer(S.k_cache, 0, 2); + enc->setBuffer(S.v_cache, 0, 3); + const uint32_t one = 1; + enc->setBytes(&one, 4, 4); + enc->setBytes(&p.n_kv_heads, 4, 5); + enc->setBytes(&hd, 4, 6); + enc->setBytes(&T, 4, 7); + enc->setBytes(&p.current_pos, 4, 8); + enc->setBytes(&p.cache_max, 4, 9); + enc->dispatchThreads(MTL::Size(hd / 4, T, p.n_kv_heads), MTL::Size(32, 4, 1)); + enc->endEncoding(); + } + + { + const uint32_t kv_len = p.current_pos + T; + const uint32_t Hg = p.n_heads / p.n_kv_heads; + const uint32_t br = 2; + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.attn); + enc->setBuffer(q_in, 0, 0); + enc->setBuffer(S.k_cache, 0, 1); + enc->setBuffer(S.v_cache, 0, 2); + enc->setBuffer(B.attn_out, 0, 3); + enc->setBytes(&T, 4, 4); + enc->setBytes(&p.n_heads, 4, 5); + enc->setBytes(&p.n_kv_heads, 4, 6); + enc->setBytes(&kv_len, 4, 7); + enc->setBytes(&p.cache_max, 4, 8); + enc->dispatchThreadgroups(MTL::Size(p.n_kv_heads, (T + br - 1) / br, 1), + MTL::Size(Hg * br * 32, 1, 1)); + enc->endEncoding(); + } + + MTL::Buffer* o_in = B.attn_out; + if (T > 1) { + enc_transpose(cmd, P.t_head_to_seq, B.attn_out, B.attn_out_seq, + T, p.n_heads, hd); + o_in = B.attn_out_seq; + } + + enc_q8_matvec(cmd, P, W.wo, o_in, B.mixer_out, T, D, qN); +} + +// ── full model ─────────────────────────────────────────────────────────────── + +inline void dispatch_model(MTL::CommandBuffer* cmd, const PSOs& P, + const Weights& W, const std::vector& S, + Buffers& B, const Params& p) +{ + const uint32_t T = p.seq; + const uint32_t D = p.d_model; + + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.embedding_lookup); + enc->setBuffer(W.embed, 0, 0); + enc->setBuffer(B.input_ids, 0, 1); + enc->setBuffer(B.x_a, 0, 2); + enc->setBytes(&T, 4, 3); + enc->setBytes(&D, 4, 4); + enc->setBytes(&p.vocab_size, 4, 5); + enc->dispatchThreadgroups(MTL::Size((D / 4 + 127) / 128, T, 1), + MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + enc_scale(cmd, P, B.x_a, p.embedding_scale, T * D); + + MTL::Buffer* cur = B.x_a; + MTL::Buffer* nxt = B.x_b; + for (uint32_t L = 0; L < p.n_layers; ++L) { + const LayerWeights& lw = W.layers[L]; + const size_t off_norm = (size_t)L * D * 2; + + // ffn_inp/mixer scratch use the layer-shared buffers; x_norm is + // recomputed per stage so cur/nxt are the only cross-layer carriers. + Buffers& b = B; + b.x_norm = B.x_norm; + + enc_rmsnorm(cmd, P, cur, W.attn_norm, off_norm, B.x_norm, T, D, p.eps); + if (lw.is_attn) dispatch_attn_mixer(cmd, P, p, lw, S[L], b); + else dispatch_mamba_mixer(cmd, P, p, lw, S[L], b); + + enc_add_scaled(cmd, P, cur, B.mixer_out, B.ffn_inp, p.residual_scale, T * D); + + enc_rmsnorm(cmd, P, B.ffn_inp, W.ffn_norm, off_norm, B.x_norm, T, D, p.eps); + enc_q8_matvec(cmd, P, lw.gate, B.x_norm, B.gate_buf, T, p.n_int, D); + enc_q8_matvec(cmd, P, lw.up, B.x_norm, B.up_buf, T, p.n_int, D); + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.silu_mul); + enc->setBuffer(B.gate_buf, 0, 0); + enc->setBuffer(B.up_buf, 0, 1); + enc->setBuffer(B.up_buf, 0, 2); + uint32_t n_total = T * p.n_int; + enc->setBytes(&n_total, 4, 3); + enc->dispatchThreadgroups(MTL::Size((n_total + 255) / 256, 1, 1), + MTL::Size(256, 1, 1)); + enc->endEncoding(); + } + enc_q8_matvec(cmd, P, lw.down, B.up_buf, B.mlp_out, T, D, p.n_int); + + enc_add_scaled(cmd, P, B.ffn_inp, B.mlp_out, nxt, p.residual_scale, T * D); + + MTL::Buffer* tmp = cur; cur = nxt; nxt = tmp; + } + + enc_rmsnorm(cmd, P, cur, W.final_norm, 0, nxt, T, D, p.eps); + + // Tied Q8_0 head on the LAST row only (prefill projects just row T-1); + // argmax is scale-invariant but the logit_scale division keeps + // get_last_logits honest. + { + const uint32_t last = T - 1; + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.q8_0_matvec); + enc->setBuffer(nxt, (size_t)last * D * 2, 0); + enc->setBuffer(W.head_q8, 0, 1); + enc->setBuffer(B.logits, 0, 2); + enc->setBytes(&D, 4, 3); + enc->setBytes(&p.vocab_size, 4, 4); + enc->dispatchThreadgroups(MTL::Size((p.vocab_size + 1) / 2, 1, 1), + MTL::Size(128, 1, 1)); + enc->endEncoding(); + } + enc_scale(cmd, P, B.logits, 1.0f / p.logit_scale, p.vocab_size); + + { + auto* enc = cmd->computeCommandEncoder(); + enc->setComputePipelineState(P.argmax); + enc->setBuffer(B.logits, 0, 0); + enc->setBuffer(B.output_id, 0, 1); + enc->setBytes(&p.vocab_size, 4, 2); + enc->dispatchThreadgroups(MTL::Size(1, 1, 1), MTL::Size(1024, 1, 1)); + enc->endEncoding(); + } +} + +} // namespace granite +} // namespace meow + +#endif diff --git a/SuperKittens/models/granite/granite_ops.metal b/SuperKittens/models/granite/granite_ops.metal new file mode 100644 index 0000000..4e2c194 --- /dev/null +++ b/SuperKittens/models/granite/granite_ops.metal @@ -0,0 +1,39 @@ +// +// granite_ops.metal — Granite-4.x hybrid elementwise ops. +// +// Granite carries four scalar multipliers (embedding 12, residual 0.22, +// attention 1/128, logits 1/6) that no shared SK kernel applies; these two +// kernels plumb them without touching shared-kernel signatures. + +#include +using namespace metal; + +// x *= s, in place. Also used to pre-scale Q so mha_causal's hardcoded +// 1/sqrt(D) softmax scale becomes granite's attention_multiplier: +// s = attention_multiplier * sqrt(D). +[[host_name("granite_scale_f16")]] +[[kernel]] +void granite_scale_f16( + device half* x [[buffer(0)]], + constant float& s [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint tid [[thread_position_in_grid]]) +{ + if (tid >= n) return; + x[tid] = half(float(x[tid]) * s); +} + +// y = a + s*b (residual_multiplier on every mixer/FFN block output). +[[host_name("granite_add_scaled_f16")]] +[[kernel]] +void granite_add_scaled_f16( + device const half* a [[buffer(0)]], + device const half* b [[buffer(1)]], + device half* y [[buffer(2)]], + constant float& s [[buffer(3)]], + constant uint& n [[buffer(4)]], + uint tid [[thread_position_in_grid]]) +{ + if (tid >= n) return; + y[tid] = half(float(a[tid]) + s * float(b[tid])); +} diff --git a/SuperKittens/models/granite/launcher.c++ b/SuperKittens/models/granite/launcher.c++ new file mode 100644 index 0000000..d1b8fd3 --- /dev/null +++ b/SuperKittens/models/granite/launcher.c++ @@ -0,0 +1,574 @@ +// launcher.c++ — Granite-4.x hybrid inference launcher. + +#include "launcher.h" +#include "granite_model.h" +#include "../../kernels/runtime_bindings.h" +#include "../load/gguf/gguf.h" + +#include +#include +#include +#include +#include +#include + +namespace meow { namespace granite { + +struct Handle { + sk_granite_config cfg; + uint32_t current_pos = 0; + uint32_t last_seq = 0; + bool loaded = false; + + PSOs psos; + Weights weights; + std::vector states; + Buffers bufs; +}; + +static MTL::Buffer* alloc_zero(MTL::Device* dev, size_t bytes) { + auto* b = dev->newBuffer(bytes, MTL::ResourceStorageModeShared); + if (b) std::memset(b->contents(), 0, bytes); + return b; +} + +static bool resolve_psos(PSOs& P) { + P.rmsnorm = sk::bindings_pso("rmsnorm"); + P.rmsnorm_t1 = sk::bindings_pso("rmsnorm_t1"); // optional + P.q8_0_matvec = sk::bindings_pso("q8_0_matvec"); + P.split_packed = sk::bindings_pso("split_packed"); + P.conv1d_silu = sk::bindings_pso("conv1d_silu"); + P.conv1d_silu_step = sk::bindings_pso("conv1d_silu_step"); + P.conv_state_capture = sk::bindings_pso("conv_state_capture"); + P.mamba2_ssd = sk::bindings_pso("mamba2_ssd"); + if (!P.mamba2_ssd) P.mamba2_ssd = sk::bindings_pso("mamba2_ssd_ref"); + P.gate_norm = sk::bindings_pso("gate_norm"); + P.silu_mul = sk::bindings_pso("silu_mul_f16"); + P.attn = sk::bindings_pso("mha_causal"); + P.kv_cache_write = sk::bindings_pso("kv_cache_write"); + P.t_seq_to_head = sk::bindings_pso("transpose_seq_to_head_f16"); + P.t_head_to_seq = sk::bindings_pso("transpose_head_to_seq_f16"); + P.scale = sk::bindings_pso("granite_scale_f16"); + P.add_scaled = sk::bindings_pso("granite_add_scaled_f16"); + P.embedding_lookup = sk::bindings_pso("embedding_lookup"); + P.argmax = sk::bindings_pso("argmax"); + + #define _CK(name, val) if (!(val)) { std::fprintf(stderr, "granite launcher: missing PSO " name "\n"); return false; } + _CK("rmsnorm", P.rmsnorm); + _CK("q8_0_matvec", P.q8_0_matvec); + _CK("split_packed", P.split_packed); + _CK("conv1d_silu", P.conv1d_silu); + _CK("conv1d_silu_step", P.conv1d_silu_step); + _CK("conv_state_capture", P.conv_state_capture); + _CK("mamba2_ssd(_ref)", P.mamba2_ssd); + _CK("gate_norm", P.gate_norm); + _CK("silu_mul_f16", P.silu_mul); + _CK("mha_causal", P.attn); + _CK("kv_cache_write", P.kv_cache_write); + _CK("transpose_seq_to_head_f16", P.t_seq_to_head); + _CK("transpose_head_to_seq_f16", P.t_head_to_seq); + _CK("granite_scale_f16", P.scale); + _CK("granite_add_scaled_f16", P.add_scaled); + _CK("embedding_lookup", P.embedding_lookup); + _CK("argmax", P.argmax); + #undef _CK + return true; +} + +// ── dtype conversion (subnormal-correct fp32->fp16; K-quant scales are tiny) ── + +static uint16_t fp32_bits_to_fp16(uint32_t f) { + uint32_t sign = (f >> 16) & 0x8000u; + uint32_t mant = f & 0x007fffffu; + int32_t exp = (int32_t)((f >> 23) & 0xffu) - 127 + 15; + if (((f >> 23) & 0xffu) == 0xffu) return (uint16_t)(sign | 0x7c00u | (mant ? 0x0200u : 0u)); + if (exp >= 31) return (uint16_t)(sign | 0x7c00u); + if (exp <= 0) { + if (exp < -10) return (uint16_t)sign; + mant = (mant | 0x00800000u) >> (uint32_t)(1 - exp); + if (mant & 0x00001000u) mant += 0x00002000u; + return (uint16_t)(sign | (mant >> 13)); + } + if (mant & 0x00001000u) { + mant += 0x00002000u; + if (mant & 0x00800000u) { mant = 0; exp += 1; } + if (exp >= 31) return (uint16_t)(sign | 0x7c00u); + } + return (uint16_t)(sign | ((uint32_t)exp << 10) | (mant >> 13)); +} + +static uint16_t f32_to_fp16(float f) { + uint32_t fb; std::memcpy(&fb, &f, 4); + return fp32_bits_to_fp16(fb); +} + +static float fp16_bits_to_f32(uint16_t s) { + uint32_t sign = (s & 0x8000u) << 16; + uint32_t exp = (s >> 10) & 0x1fu; + uint32_t mant = s & 0x3ffu; + uint32_t v; + if (exp == 0) { + if (mant == 0) v = sign; + else { + exp = 1; + while ((mant & 0x400u) == 0) { mant <<= 1; exp -= 1; } + mant &= 0x3ffu; + v = sign | ((exp + 112) << 23) | (mant << 13); + } + } else if (exp == 31) { + v = sign | 0x7f800000u | (mant << 13); + } else { + v = sign | ((exp + 112) << 23) | (mant << 13); + } + float out; std::memcpy(&out, &v, 4); + return out; +} + +static void dequant_q8_0_to_fp16(uint16_t* dst, const uint8_t* src, size_t n_elems) { + const size_t n_blocks = n_elems / 32; + for (size_t b = 0; b < n_blocks; ++b) { + const uint8_t* p = src + b * 34; + uint16_t scale_h; std::memcpy(&scale_h, p, 2); + const float scale = fp16_bits_to_f32(scale_h); + const int8_t* qs = (const int8_t*)(p + 2); + for (int i = 0; i < 32; ++i) + dst[b * 32 + i] = f32_to_fp16((float)qs[i] * scale); + } +} + +// ── GGUF load ───────────────────────────────────────────────────────────────── + +static bool copy_f32_as_fp16(MTL::Buffer* dst, size_t dst_off, + sk::WeightStore& store, const std::string& name, + size_t n_elems) { + auto* v = store.get(name); + if (!v) { std::fprintf(stderr, "granite gguf: missing %s\n", name.c_str()); return false; } + if (v->dtype != sk::Dtype::F32 || v->nbytes != n_elems * 4) { + std::fprintf(stderr, "granite gguf: %s dtype/size unexpected (dtype=%d nbytes=%zu want F32 x%zu)\n", + name.c_str(), (int)v->dtype, v->nbytes, n_elems); + return false; + } + const float* s = (const float*)v->data; + uint16_t* d = (uint16_t*)((char*)dst->contents() + dst_off); + for (size_t i = 0; i < n_elems; ++i) d[i] = f32_to_fp16(s[i]); + return true; +} + +static MTL::Buffer* copy_q8_0(MTL::Device* dev, sk::WeightStore& store, + const std::string& name, size_t n_elems) { + auto* v = store.get(name); + if (!v) { std::fprintf(stderr, "granite gguf: missing %s\n", name.c_str()); return nullptr; } + const size_t expect = (n_elems / 32) * 34; + if (v->dtype != sk::Dtype::Q8_0 || v->nbytes != expect) { + std::fprintf(stderr, "granite gguf: %s dtype/size unexpected (dtype=%d nbytes=%zu want Q8_0 %zu)\n", + name.c_str(), (int)v->dtype, v->nbytes, expect); + return nullptr; + } + auto* b = dev->newBuffer(expect, MTL::ResourceStorageModeShared); + if (!b) return nullptr; + std::memcpy(b->contents(), v->data, expect); + return b; +} + +}} // namespace meow::granite + +extern "C" sk_granite_handle* sk_granite_create(const sk_granite_config* cfg) { + if (!cfg) return nullptr; + auto* dev = sk::bindings_device(); + if (!dev) return nullptr; + + auto* h = new meow::granite::Handle(); + h->cfg = *cfg; + if (!meow::granite::resolve_psos(h->psos)) { delete h; return nullptr; } + + using namespace meow::granite; + const uint32_t T_max = cfg->seq_max; + const uint32_t D = cfg->d_model; + const uint32_t E = cfg->d_inner; + const uint32_t C_in = E + 2 * cfg->ssm_n_groups * cfg->ssm_state; + const uint32_t IN_OUT = 2 * E + 2 * cfg->ssm_n_groups * cfg->ssm_state + cfg->ssm_n_heads; + const uint32_t qN = cfg->n_heads * cfg->head_dim; + const uint32_t kvN = cfg->n_kv_heads * cfg->head_dim; + + h->weights.embed = alloc_zero(dev, (size_t)cfg->vocab_size * D * 2); + h->weights.attn_norm = alloc_zero(dev, (size_t)cfg->n_layers * D * 2); + h->weights.ffn_norm = alloc_zero(dev, (size_t)cfg->n_layers * D * 2); + h->weights.final_norm = alloc_zero(dev, (size_t)D * 2); + + auto& b = h->bufs; + b.input_ids = alloc_zero(dev, (size_t)T_max * sizeof(int32_t)); + b.output_id = alloc_zero(dev, sizeof(int32_t)); + b.x_a = alloc_zero(dev, (size_t)T_max * D * 2); + b.x_b = alloc_zero(dev, (size_t)T_max * D * 2); + b.x_norm = alloc_zero(dev, (size_t)T_max * D * 2); + b.ffn_inp = alloc_zero(dev, (size_t)T_max * D * 2); + b.mixer_out = alloc_zero(dev, (size_t)T_max * D * 2); + b.logits = alloc_zero(dev, (size_t)cfg->vocab_size * 2); + b.in_proj_out = alloc_zero(dev, (size_t)T_max * IN_OUT * 2); + b.z = alloc_zero(dev, (size_t)T_max * E * 2); + b.xBC = alloc_zero(dev, (size_t)T_max * C_in * 2); + b.dt_raw = alloc_zero(dev, (size_t)T_max * cfg->ssm_n_heads * 2); + b.xBC_post = alloc_zero(dev, (size_t)T_max * (C_in + cfg->ssm_n_heads) * 2); + b.ssd_out = alloc_zero(dev, (size_t)T_max * E * 2); + b.gated = alloc_zero(dev, (size_t)T_max * E * 2); + b.q = alloc_zero(dev, (size_t)T_max * qN * 2); + b.k_tmp = alloc_zero(dev, (size_t)T_max * kvN * 2); + b.v_tmp = alloc_zero(dev, (size_t)T_max * kvN * 2); + b.attn_out = alloc_zero(dev, (size_t)T_max * qN * 2); + b.q_th = alloc_zero(dev, (size_t)T_max * qN * 2); + b.k_th = alloc_zero(dev, (size_t)T_max * kvN * 2); + b.v_th = alloc_zero(dev, (size_t)T_max * kvN * 2); + b.attn_out_seq = alloc_zero(dev, (size_t)T_max * qN * 2); + b.gate_buf = alloc_zero(dev, (size_t)T_max * cfg->n_int * 2); + b.up_buf = alloc_zero(dev, (size_t)T_max * cfg->n_int * 2); + b.mlp_out = alloc_zero(dev, (size_t)T_max * D * 2); + + return reinterpret_cast(h); +} + +extern "C" int sk_granite_load_gguf(sk_granite_handle* hp, const char* path) { + if (!hp || !path) return -1; + auto* h = reinterpret_cast(hp); + auto* dev = sk::bindings_device(); + if (!dev) return -2; + const auto& c = h->cfg; + + sk::WeightStore store; + sk::gguf::Model gm; + int rc = sk::gguf::load_gguf(path, store, &gm); + if (rc != 0) { + std::fprintf(stderr, "granite gguf: parse failed rc=%d\n", rc); + return rc; + } + + std::string arch; + sk::gguf::meta_string(gm, "general.architecture", &arch); + if (arch != "granitehybrid") { + std::fprintf(stderr, "granite gguf: arch '%s' != granitehybrid\n", arch.c_str()); + return -10; + } + + // Layer types: head_count_kv is a per-layer U32 array; 0 = mamba (recurrent). + const auto* kvh = sk::gguf::meta_find(gm, "granitehybrid.attention.head_count_kv"); + if (!kvh || kvh->type != sk::gguf::V_ARRAY || kvh->arr_type != sk::gguf::V_U32 + || kvh->arr_len != c.n_layers) { + std::fprintf(stderr, "granite gguf: head_count_kv array missing/mismatched\n"); + return -11; + } + std::vector kv_heads(c.n_layers); + std::memcpy(kv_heads.data(), gm.map_base + kvh->raw_pos, (size_t)c.n_layers * 4); + + // Dim cross-check against metadata (fail loud on mismatch, not garbage). + auto check_u32 = [&](const char* key, uint32_t want) -> bool { + uint32_t v = 0; + if (!sk::gguf::meta_u32(gm, key, &v)) { + std::fprintf(stderr, "granite gguf: missing meta %s\n", key); + return false; + } + if (v != want) { + std::fprintf(stderr, "granite gguf: meta %s=%u != cfg %u\n", key, v, want); + return false; + } + return true; + }; + if (!check_u32("granitehybrid.block_count", c.n_layers)) return -12; + if (!check_u32("granitehybrid.embedding_length", c.d_model)) return -12; + if (!check_u32("granitehybrid.feed_forward_length", c.n_int)) return -12; + if (!check_u32("granitehybrid.attention.head_count", c.n_heads)) return -12; + if (!check_u32("granitehybrid.vocab_size", c.vocab_size)) return -12; + if (!check_u32("granitehybrid.ssm.inner_size", c.d_inner)) return -12; + if (!check_u32("granitehybrid.ssm.state_size", c.ssm_state)) return -12; + if (!check_u32("granitehybrid.ssm.time_step_rank", c.ssm_n_heads))return -12; + if (!check_u32("granitehybrid.ssm.group_count", c.ssm_n_groups))return -12; + if (!check_u32("granitehybrid.ssm.conv_kernel", c.ssm_conv)) return -12; + + float m_attn = 0, m_embd = 0, m_res = 0, m_logit = 0; + sk::gguf::meta_f32(gm, "granitehybrid.attention.scale", &m_attn); + sk::gguf::meta_f32(gm, "granitehybrid.embedding_scale", &m_embd); + sk::gguf::meta_f32(gm, "granitehybrid.residual_scale", &m_res); + sk::gguf::meta_f32(gm, "granitehybrid.logit_scale", &m_logit); + + // Gate A config table. + std::fprintf(stderr, "granite config: layers=%u d_model=%u ffn=%u vocab=%u\n", + c.n_layers, c.d_model, c.n_int, c.vocab_size); + std::fprintf(stderr, "granite attn: heads=%u kv_heads=%u head_dim=%u NoPE scale=%g\n", + c.n_heads, c.n_kv_heads, c.head_dim, m_attn); + std::fprintf(stderr, "granite ssm: E=%u H=%u P=%u N=%u G=%u K=%u\n", + c.d_inner, c.ssm_n_heads, c.ssm_head_dim, c.ssm_state, + c.ssm_n_groups, c.ssm_conv); + std::fprintf(stderr, "granite scales: embed=%g residual=%g attn=%g logit=%g\n", + m_embd, m_res, m_attn, m_logit); + std::fprintf(stderr, "granite layer types: "); + for (uint32_t L = 0; L < c.n_layers; ++L) + std::fprintf(stderr, "%c", kv_heads[L] > 0 ? 'A' : 'm'); + std::fprintf(stderr, " (A=attention, m=mamba2)\n"); + + using namespace meow::granite; + const uint32_t D = c.d_model; + const uint32_t E = c.d_inner; + const uint32_t H = c.ssm_n_heads; + const uint32_t G = c.ssm_n_groups; + const uint32_t N = c.ssm_state; + const uint32_t K = c.ssm_conv; + const uint32_t C_in = E + 2 * G * N; + const uint32_t IN_OUT = 2 * E + 2 * G * N + H; + const uint32_t qN = c.n_heads * c.head_dim; + const uint32_t kvN = c.n_kv_heads * c.head_dim; + + // Embedding: fp16 dequant for the lookup; raw Q8_0 copy for the tied head. + { + auto* v = store.get("token_embd.weight"); + if (!v) { std::fprintf(stderr, "granite gguf: missing token_embd.weight\n"); return -20; } + const size_t n_elems = (size_t)c.vocab_size * D; + if (v->dtype != sk::Dtype::Q8_0 || v->nbytes != (n_elems / 32) * 34) { + std::fprintf(stderr, "granite gguf: token_embd not Q8_0 (dtype=%d)\n", (int)v->dtype); + return -20; + } + dequant_q8_0_to_fp16((uint16_t*)h->weights.embed->contents(), + (const uint8_t*)v->data, n_elems); + h->weights.head_q8 = copy_q8_0(dev, store, "token_embd.weight", n_elems); + if (!h->weights.head_q8) return -21; + } + if (!copy_f32_as_fp16(h->weights.final_norm, 0, store, "output_norm.weight", D)) + return -22; + + h->weights.layers.assign(c.n_layers, LayerWeights{}); + h->states.assign(c.n_layers, LayerState{}); + + char nm[128]; + for (uint32_t L = 0; L < c.n_layers; ++L) { + LayerWeights& lw = h->weights.layers[L]; + LayerState& ls = h->states[L]; + lw.is_attn = kv_heads[L] > 0; + if (lw.is_attn && kv_heads[L] != c.n_kv_heads) { + std::fprintf(stderr, "granite gguf: layer %u kv_heads=%u != cfg %u\n", + L, kv_heads[L], c.n_kv_heads); + return -23; + } + + const size_t off_norm = (size_t)L * D * 2; + std::snprintf(nm, sizeof(nm), "blk.%u.attn_norm.weight", L); + if (!copy_f32_as_fp16(h->weights.attn_norm, off_norm, store, nm, D)) return -30; + std::snprintf(nm, sizeof(nm), "blk.%u.ffn_norm.weight", L); + if (!copy_f32_as_fp16(h->weights.ffn_norm, off_norm, store, nm, D)) return -31; + + if (lw.is_attn) { + std::snprintf(nm, sizeof(nm), "blk.%u.attn_q.weight", L); + lw.wq = copy_q8_0(dev, store, nm, (size_t)qN * D); + std::snprintf(nm, sizeof(nm), "blk.%u.attn_k.weight", L); + lw.wk = copy_q8_0(dev, store, nm, (size_t)kvN * D); + std::snprintf(nm, sizeof(nm), "blk.%u.attn_v.weight", L); + lw.wv = copy_q8_0(dev, store, nm, (size_t)kvN * D); + std::snprintf(nm, sizeof(nm), "blk.%u.attn_output.weight", L); + lw.wo = copy_q8_0(dev, store, nm, (size_t)D * qN); + if (!lw.wq || !lw.wk || !lw.wv || !lw.wo) return -40; + + ls.k_cache = alloc_zero(dev, (size_t)c.cache_max * kvN * 2); + ls.v_cache = alloc_zero(dev, (size_t)c.cache_max * kvN * 2); + } else { + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_in.weight", L); + lw.ssm_in = copy_q8_0(dev, store, nm, (size_t)IN_OUT * D); + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_out.weight", L); + lw.ssm_out = copy_q8_0(dev, store, nm, (size_t)D * E); + if (!lw.ssm_in || !lw.ssm_out) return -41; + + // conv1d weight: GGUF layout [C_in rows][K] matches conv1d_silu's + // w[c*K+k] indexing — straight F32->fp16, no transpose. + lw.conv_w = alloc_zero(dev, (size_t)C_in * K * 2); + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_conv1d.weight", L); + if (!copy_f32_as_fp16(lw.conv_w, 0, store, nm, (size_t)C_in * K)) return -42; + lw.conv_b = alloc_zero(dev, (size_t)C_in * 2); + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_conv1d.bias", L); + if (!copy_f32_as_fp16(lw.conv_b, 0, store, nm, C_in)) return -43; + + lw.dt_bias = alloc_zero(dev, (size_t)H * 2); + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_dt.bias", L); + if (!copy_f32_as_fp16(lw.dt_bias, 0, store, nm, H)) return -44; + + // GGUF stores A = -exp(A_log) (llama.cpp conversion); SK's SSD + // kernels compute dA = exp(dt * -exp(A_log)), so invert here. + { + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_a", L); + auto* v = store.get(nm); + if (!v || v->dtype != sk::Dtype::F32 || v->nbytes != (size_t)H * 4) { + std::fprintf(stderr, "granite gguf: %s missing/unexpected\n", nm); + return -45; + } + lw.A_log = alloc_zero(dev, (size_t)H * 2); + const float* a = (const float*)v->data; + uint16_t* d = (uint16_t*)lw.A_log->contents(); + for (uint32_t i = 0; i < H; ++i) { + if (!(a[i] < 0.0f)) { + std::fprintf(stderr, "granite gguf: %s[%u]=%g not negative\n", nm, i, a[i]); + return -45; + } + d[i] = f32_to_fp16(std::log(-a[i])); + } + } + + lw.ssm_D = alloc_zero(dev, (size_t)H * 2); + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_d", L); + if (!copy_f32_as_fp16(lw.ssm_D, 0, store, nm, H)) return -46; + lw.ssm_norm = alloc_zero(dev, (size_t)E * 2); + std::snprintf(nm, sizeof(nm), "blk.%u.ssm_norm.weight", L); + if (!copy_f32_as_fp16(lw.ssm_norm, 0, store, nm, E)) return -47; + + ls.conv_state = alloc_zero(dev, (size_t)(K - 1) * C_in * 2); + ls.ssm_state = alloc_zero(dev, (size_t)H * c.ssm_head_dim * N * sizeof(float)); + } + + std::snprintf(nm, sizeof(nm), "blk.%u.ffn_gate.weight", L); + lw.gate = copy_q8_0(dev, store, nm, (size_t)c.n_int * D); + std::snprintf(nm, sizeof(nm), "blk.%u.ffn_up.weight", L); + lw.up = copy_q8_0(dev, store, nm, (size_t)c.n_int * D); + std::snprintf(nm, sizeof(nm), "blk.%u.ffn_down.weight", L); + lw.down = copy_q8_0(dev, store, nm, (size_t)D * c.n_int); + if (!lw.gate || !lw.up || !lw.down) return -48; + } + + h->loaded = true; + return 0; +} + +namespace meow { namespace granite { + +static int run_step(Handle* h, MTL::CommandQueue* q, uint32_t seq) { + Params p; + p.seq = seq; + p.n_layers = h->cfg.n_layers; + p.d_model = h->cfg.d_model; + p.n_heads = h->cfg.n_heads; + p.n_kv_heads = h->cfg.n_kv_heads; + p.head_dim = h->cfg.head_dim; + p.n_int = h->cfg.n_int; + p.d_inner = h->cfg.d_inner; + p.ssm_heads = h->cfg.ssm_n_heads; + p.ssm_pdim = h->cfg.ssm_head_dim; + p.ssm_state = h->cfg.ssm_state; + p.ssm_groups = h->cfg.ssm_n_groups; + p.ssm_conv = h->cfg.ssm_conv; + p.vocab_size = h->cfg.vocab_size; + p.cache_max = h->cfg.cache_max; + p.current_pos = h->current_pos; + p.eps = h->cfg.eps; + p.embedding_scale = h->cfg.embedding_scale; + p.residual_scale = h->cfg.residual_scale; + p.attention_scale = h->cfg.attention_scale; + p.logit_scale = h->cfg.logit_scale; + h->last_seq = seq; + + auto* cmd = q->commandBuffer(); + dispatch_model(cmd, h->psos, h->weights, h->states, h->bufs, p); + cmd->commit(); + cmd->waitUntilCompleted(); + if (getenv("SK_GRANITE_GPUPROF")) + std::fprintf(stderr, "[gpuprof] gpu_busy_us=%.1f\n", + (cmd->GPUEndTime() - cmd->GPUStartTime()) * 1e6); + cmd->release(); + h->current_pos += seq; + return 0; +} + +}} // namespace meow::granite + +extern "C" int sk_granite_forward(sk_granite_handle* hp, + const int* input_ids, uint32_t seq, int* output_id) { + if (!hp || !input_ids || !output_id) return -1; + auto* h = reinterpret_cast(hp); + if (!h->loaded) return -7; + if (seq == 0 || seq > h->cfg.seq_max) return -2; + // Chunked mamba prefill is unsupported (each chunk would re-zero-pad the + // conv left edge); a prompt must fit one prefill forward. + if (h->current_pos > 0 && seq > 1) return -6; + if (h->current_pos + seq > h->cfg.cache_max) return -4; + auto* q = sk::bindings_queue(); + if (!q) return -3; + + std::memcpy(h->bufs.input_ids->contents(), input_ids, (size_t)seq * sizeof(int32_t)); + if (int rc = meow::granite::run_step(h, q, seq)) return rc; + std::memcpy(output_id, h->bufs.output_id->contents(), sizeof(int32_t)); + return 0; +} + +extern "C" int sk_granite_generate_n(sk_granite_handle* hp, + const int* prompt_ids, uint32_t prompt_seq, + int* out_tokens, uint32_t n_tokens, int32_t eos_id) { + if (!hp || !prompt_ids || !out_tokens) return -1; + auto* h = reinterpret_cast(hp); + if (!h->loaded) return -7; + if (prompt_seq == 0 || prompt_seq > h->cfg.seq_max) return -2; + if (n_tokens == 0) return 0; + if (h->current_pos != 0) return -6; // mamba state demands a fresh sequence + auto* q = sk::bindings_queue(); + if (!q) return -3; + + std::memcpy(h->bufs.input_ids->contents(), prompt_ids, + (size_t)prompt_seq * sizeof(int32_t)); + if (int rc = meow::granite::run_step(h, q, prompt_seq)) return rc; + + int32_t* in_ids = (int32_t*)h->bufs.input_ids->contents(); + const int32_t* out_id = (const int32_t*)h->bufs.output_id->contents(); + + int32_t last = out_id[0]; + out_tokens[0] = last; + if (eos_id >= 0 && last == eos_id) return 1; + uint32_t written = 1; + while (written < n_tokens) { + if (h->current_pos + 1 > h->cfg.cache_max) break; + in_ids[0] = last; + if (int rc = meow::granite::run_step(h, q, 1)) return rc; + last = out_id[0]; + out_tokens[written++] = last; + if (eos_id >= 0 && last == eos_id) break; + } + return (int)written; +} + +extern "C" int sk_granite_get_last_logits(sk_granite_handle* hp, void* out_fp16) { + if (!hp || !out_fp16) return -1; + auto* h = reinterpret_cast(hp); + if (h->last_seq == 0) return -2; + std::memcpy(out_fp16, h->bufs.logits->contents(), + (size_t)h->cfg.vocab_size * 2); + return 0; +} + +extern "C" uint32_t sk_granite_get_pos(sk_granite_handle* hp) { + if (!hp) return 0; + return reinterpret_cast(hp)->current_pos; +} + +extern "C" void sk_granite_reset(sk_granite_handle* hp) { + if (!hp) return; + auto* h = reinterpret_cast(hp); + h->current_pos = 0; + for (auto& s : h->states) { + if (s.conv_state) std::memset(s.conv_state->contents(), 0, s.conv_state->length()); + if (s.ssm_state) std::memset(s.ssm_state->contents(), 0, s.ssm_state->length()); + } +} + +extern "C" void sk_granite_destroy(sk_granite_handle* hp) { + if (!hp) return; + auto* h = reinterpret_cast(hp); + auto rel = [](MTL::Buffer* b) { if (b) b->release(); }; + rel(h->weights.embed); rel(h->weights.head_q8); + rel(h->weights.attn_norm); rel(h->weights.ffn_norm); rel(h->weights.final_norm); + for (auto& lw : h->weights.layers) { + rel(lw.ssm_in); rel(lw.ssm_out); rel(lw.conv_w); rel(lw.conv_b); + rel(lw.dt_bias); rel(lw.A_log); rel(lw.ssm_D); rel(lw.ssm_norm); + rel(lw.wq); rel(lw.wk); rel(lw.wv); rel(lw.wo); + rel(lw.gate); rel(lw.up); rel(lw.down); + } + for (auto& s : h->states) { + rel(s.k_cache); rel(s.v_cache); rel(s.conv_state); rel(s.ssm_state); + } + auto& b = h->bufs; + rel(b.input_ids); rel(b.output_id); rel(b.x_a); rel(b.x_b); rel(b.x_norm); + rel(b.ffn_inp); rel(b.mixer_out); rel(b.logits); rel(b.in_proj_out); + rel(b.z); rel(b.xBC); rel(b.dt_raw); rel(b.xBC_post); rel(b.ssd_out); + rel(b.gated); rel(b.q); rel(b.k_tmp); rel(b.v_tmp); rel(b.attn_out); + rel(b.q_th); rel(b.k_th); rel(b.v_th); rel(b.attn_out_seq); + rel(b.gate_buf); rel(b.up_buf); rel(b.mlp_out); + delete h; +} diff --git a/SuperKittens/models/granite/launcher.h b/SuperKittens/models/granite/launcher.h new file mode 100644 index 0000000..5777bf0 --- /dev/null +++ b/SuperKittens/models/granite/launcher.h @@ -0,0 +1,78 @@ +// +// launcher.h — IBM Granite-4.x hybrid (interleaved mamba2 + attention) C ABI. +// +// Layer types come from GGUF metadata at load time +// (granitehybrid.attention.head_count_kv: per-layer kv-head count, 0 = mamba). +// Every layer additionally carries a dense SwiGLU FFN. + +#ifndef SK_GRANITE_LAUNCHER_H +#define SK_GRANITE_LAUNCHER_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct { + uint32_t batch; // 1 (single-stream stage 1) + uint32_t seq_max; // max tokens per forward (prompt must fit ONE prefill: + // chunked mamba prefill would zero-pad each chunk's conv left edge) + uint32_t cache_max; // attention KV capacity + uint32_t n_layers; // 40 + uint32_t d_model; // 1536 + // attention layers + uint32_t n_heads; // 12 + uint32_t n_kv_heads; // 4 + uint32_t head_dim; // 128 (mha_causal D=128 instantiation) + // dense FFN (every layer) + uint32_t n_int; // 4096 + // mamba2 layers + uint32_t d_inner; // 3072 (E) + uint32_t ssm_n_heads; // 48 (H) + uint32_t ssm_head_dim; // 64 (P) + uint32_t ssm_state; // 128 (N) + uint32_t ssm_n_groups; // 1 (G) + uint32_t ssm_conv; // 4 (K) + uint32_t vocab_size; // 100352 + float eps; // 1e-5 + // granite multipliers (GGUF: granitehybrid.{embedding,residual,attention,logit}_scale) + float embedding_scale; // 12.0 + float residual_scale; // 0.22 + float attention_scale; // 0.0078125 (replaces 1/sqrt(head_dim)) + float logit_scale; // 6.0 (logits are DIVIDED by this) +} sk_granite_config; + +typedef struct sk_granite_handle sk_granite_handle; + +sk_granite_handle* sk_granite_create(const sk_granite_config* cfg); + +// Loads weights AND the per-layer type list from GGUF metadata. Q8_0 + F32 +// GGUFs only (granite official Q8_0). +int sk_granite_load_gguf(sk_granite_handle* h, const char* path); + +// One forward (prefill seq>1 or decode seq==1); greedy argmax -> *output_id. +int sk_granite_forward(sk_granite_handle* h, + const int* input_ids, uint32_t seq, int* output_id); + +// Greedy loop in C: prefill prompt then decode n_tokens (stops at eos_id >= 0). +// Returns number of tokens written, or negative error. +int sk_granite_generate_n(sk_granite_handle* h, + const int* prompt_ids, uint32_t prompt_seq, + int* out_tokens, uint32_t n_tokens, int32_t eos_id); + +// fp16 logits of the last projected row (vocab_size entries, post logit_scale). +int sk_granite_get_last_logits(sk_granite_handle* h, void* out_fp16); + +uint32_t sk_granite_get_pos(sk_granite_handle* h); + +// Zero attention position + all mamba conv/ssm state. +void sk_granite_reset(sk_granite_handle* h); + +void sk_granite_destroy(sk_granite_handle* h); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/SuperKittens/models/load/tokenizer/tokenizer.py b/SuperKittens/models/load/tokenizer/tokenizer.py index e055d63..685598d 100644 --- a/SuperKittens/models/load/tokenizer/tokenizer.py +++ b/SuperKittens/models/load/tokenizer/tokenizer.py @@ -60,6 +60,11 @@ class Tokenizer: "deepseek": {"bos": ("<｜begin▁of▁sentence｜>",), "eos": ("<｜end▁of▁sentence｜>",), "pad": ("<｜end▁of▁sentence｜>",)}, + # Granite 4.x: GPT-2-lineage BPE (dbrx pre-tokenizer); <|end_of_text|> + # (100257) is BOS and EOS both (GGUF bos_token_id == eos_token_id), + # add_bos=False so generation prompts are raw completions. + "granite": {"bos": ("<|end_of_text|>",), "eos": ("<|end_of_text|>",), + "pad": ("<|pad|>", "<|end_of_text|>")}, } _DEFAULT_SPECIALS: ClassVar[dict] = { "bos": ("<|im_start|>", "", "", "<|startoftext|>"), From bcb5b6137dbfa557cc6a11f2bd4fe11195d4f526 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 00:35:22 -0400 Subject: [PATCH 24/31] =?UTF-8?q?diffgemma=20stage2:=20persistent=20activa?= =?UTF-8?q?tion=20scratch=20(gemm/attn/packed-MoE=20buffers)=20+=20early?= =?UTF-8?q?=20probs16=20+=20in-place=20softcap=20=E2=80=94=20multi-forward?= =?UTF-8?q?=20generation=20swap-stormed=20the=20tight-disk=2016GB=20host?= =?UTF-8?q?=20(swapfiles=20ate=20root=20disk,=20watchdog/jetsam=20kills);?= =?UTF-8?q?=20forward=20verified=20BIT-IDENTICAL=20pre/post=20(zero-SC=20a?= =?UTF-8?q?nd=20SC=20legs)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../models/gemma/diffusion/forward_metal.py | 128 +++++++++++------- .../models/gemma/diffusion/generate.py | 31 ++++- 2 files changed, 107 insertions(+), 52 deletions(-) diff --git a/SuperKittens/models/gemma/diffusion/forward_metal.py b/SuperKittens/models/gemma/diffusion/forward_metal.py index 6debffb..9543d30 100644 --- a/SuperKittens/models/gemma/diffusion/forward_metal.py +++ b/SuperKittens/models/gemma/diffusion/forward_metal.py @@ -60,6 +60,24 @@ def __init__(self): raise RuntimeError(f"metal compile failed: {err}") self.lib = lib self._pso = {} + self._scratch = {} + + def scratch(self, tag: str, nbytes: int): + """Persistent grow-only buffer per call-site tag. Per-call Metal + alloc/free churn (activations, ~3 GB/forward) swap-stormed the 16 GB + host once generation ran many forwards in one process — same failure + class Stage 1 hit with weights, same fix.""" + b = self._scratch.get(tag) + if b is None or b.length() < nbytes: + b = self.buf_empty(nbytes) + self._scratch[tag] = b + return b + + def scratch_fill(self, tag: str, arr: np.ndarray): + arr = np.ascontiguousarray(arr) + b = self.scratch(tag, arr.nbytes) + b.contents().as_buffer(arr.nbytes)[:] = arr.tobytes() + return b def pso(self, name: str): p = self._pso.get(name) @@ -255,17 +273,21 @@ def _sc_soft_embed(self, probs16: np.ndarray) -> np.ndarray: pass return self.ctx.read(o_buf, np.float32, (C, d)) - def _sc_signal(self, sc_logits: np.ndarray, sc_temp_inv: float, - sc_use: float) -> np.ndarray: - """PR dg_canvas_embed SC subgraph -> sc_sig f32 [C, d_model].""" + def _sc_signal(self, sc_logits: np.ndarray | None, sc_temp_inv: float, + sc_use: float, probs16: np.ndarray | None = None) -> np.ndarray: + """PR dg_canvas_embed SC subgraph -> sc_sig f32 [C, d_model]. + probs16 (precomputed softmax(prev/t) f16) lets the generation loop + free the 268 MB raw-logit block between forwards.""" cfg, w = self.cfg, self.w - C, V = sc_logits.shape - # softmax(prev raw logits / prev t), fp16 on the wire like the - # reference (ggml converts the f32 probs to the f16 vec_dot type) - probs16 = np.empty((C, V), np.float16) - for c0 in range(0, C, 32): - c1 = min(c0 + 32, C) - probs16[c0:c1] = softmax(sc_logits[c0:c1] * F32(sc_temp_inv)).astype(np.float16) + if probs16 is None: + assert sc_logits is not None + C, V = sc_logits.shape + # softmax(prev raw logits / prev t), fp16 on the wire like the + # reference (ggml converts the f32 probs to the f16 vec_dot type) + probs16 = np.empty((C, V), np.float16) + for c0 in range(0, C, 32): + c1 = min(c0 + 32, C) + probs16[c0:c1] = softmax(sc_logits[c0:c1] * F32(sc_temp_inv)).astype(np.float16) if self.sc_embt_path is not None: soft = self._sc_soft_embed(probs16) else: # oracle-style host fallback (slow: full embed dequant per step) @@ -299,8 +321,8 @@ def _gemm_f32(self, wname: str, a: np.ndarray, N: int, row0: int = 0) -> np.ndar amax = float(np.abs(a).max(initial=0.0)) if amax > 3.0e4: print(f"[warn] fp16 activation near overflow ({amax:.1f}) into {wname}") - a_buf = self.ctx.buf_from(a.astype(np.float16)) - c_buf = self.ctx.buf_empty(M * N * 2) + a_buf = self.ctx.scratch_fill(f"A:{wname}", a.astype(np.float16)) + c_buf = self.ctx.scratch(f"C:{wname}", M * N * 2) b = GemmBatch(self.ctx) self._wgemm(b, wname, a_buf, 0, c_buf, 0, M, N, K, row0=row0) b.run() @@ -323,13 +345,13 @@ def _attention(self, il: int, q: np.ndarray, k: np.ndarray, v: np.ndarray, vtp = np.zeros((Kv, hd, Np), np.float16) vtp[:, :, :N] = v.transpose(1, 2, 0) - q_buf = self.ctx.buf_from(qp) - k_buf = self.ctx.buf_from(kp) - vt_buf = self.ctx.buf_from(vtp) - m_buf = self.ctx.buf_from(np.ascontiguousarray(mask[:, :Np])) - s_buf = self.ctx.buf_empty(H * N * Np * 4) # f32 scores (kq needs range) - p_buf = self.ctx.buf_empty(H * N * Np * 2) # f16 probs - o_buf = self.ctx.buf_empty(H * N * hd * 2) + q_buf = self.ctx.scratch_fill("attn_q", qp) + k_buf = self.ctx.scratch_fill("attn_k", kp) + vt_buf = self.ctx.scratch_fill("attn_vt", vtp) + m_buf = self.ctx.scratch_fill("attn_mask", np.ascontiguousarray(mask[:, :Np])) + s_buf = self.ctx.scratch("attn_s", H * N * Np * 4) # f32 scores (kq needs range) + p_buf = self.ctx.scratch("attn_p", H * N * Np * 2) # f16 probs + o_buf = self.ctx.scratch("attn_o", H * N * hd * 2) b = GemmBatch(self.ctx) for h in range(H): @@ -362,47 +384,55 @@ def _moe(self, il: int, attn_out: np.ndarray, e_in: np.ndarray) -> np.ndarray: gu_name = f"blk.{il}.ffn_gate_up_exps.weight" dn_name = f"blk.{il}.ffn_down_exps.weight" + d, nff = cfg.d_model, cfg.n_ff_exp experts = [] + r0 = 0 for e in np.unique(sel): tok, slot = np.nonzero(sel == e) - experts.append((int(e), tok, slot)) + experts.append((int(e), tok, slot, r0)) + r0 += len(tok) + m_total = r0 + + # All hit experts share three persistent row-packed scratch buffers + # (per-expert alloc/free churned ~GBs of Metal allocations per + # forward; fine for one Stage-1 forward, deadly across a generation) + cap_rows = N * cfg.n_expert_used + a_buf = self.ctx.scratch("moe_a", cap_rows * d * 2) + c_buf = self.ctx.scratch("moe_c", cap_rows * 2 * nff * 2) + d_buf = self.ctx.scratch("moe_d", cap_rows * d * 2) + packed = np.concatenate([e_in[tok] for _, tok, _, _ in experts]).astype(np.float16) + a_buf.contents().as_buffer(packed.nbytes)[:] = packed.tobytes() # stage A: gate_up for every hit expert in one command buffer b = GemmBatch(self.ctx) - stage = [] - for e, tok, slot in experts: - m = len(tok) - a_buf = self.ctx.buf_from(e_in[tok].astype(np.float16)) - c_buf = self.ctx.buf_empty(m * 2 * cfg.n_ff_exp * 2) - self._wgemm(b, gu_name, a_buf, 0, c_buf, 0, M=m, N=2 * cfg.n_ff_exp, - K=cfg.d_model, row0=e * 2 * cfg.n_ff_exp) - stage.append((e, tok, slot, c_buf, m)) + for e, tok, _, r0 in experts: + self._wgemm(b, gu_name, a_buf, r0 * d * 2, c_buf, r0 * 2 * nff * 2, + M=len(tok), N=2 * nff, K=d, row0=e * 2 * nff) b.run() - # host geglu, then stage B: down for every hit expert + # host geglu (whole packed block), then stage B: down per expert + gu = self.ctx.read(c_buf, np.float16, (m_total, 2 * nff)).astype(F32) + act = (gelu_tanh(gu[:, :nff]) * gu[:, nff:]).astype(np.float16) + a_buf.contents().as_buffer(act.nbytes)[:] = act.tobytes() b = GemmBatch(self.ctx) - stage2 = [] - for e, tok, slot, c_buf, m in stage: - gu = self.ctx.read(c_buf, np.float16, (m, 2 * cfg.n_ff_exp)).astype(F32) - act = gelu_tanh(gu[:, :cfg.n_ff_exp]) * gu[:, cfg.n_ff_exp:] - a_buf = self.ctx.buf_from(act.astype(np.float16)) - d_buf = self.ctx.buf_empty(m * cfg.d_model * 2) - self._wgemm(b, dn_name, a_buf, 0, d_buf, 0, M=m, N=cfg.d_model, - K=cfg.n_ff_exp, row0=e * cfg.d_model) - stage2.append((e, tok, slot, d_buf, m)) + for e, tok, _, r0 in experts: + self._wgemm(b, dn_name, a_buf, r0 * nff * 2, d_buf, r0 * d * 2, + M=len(tok), N=d, K=nff, row0=e * d) b.run() - moe = np.zeros((N, cfg.d_model), dtype=F32) - for e, tok, slot, d_buf, m in stage2: - d_ = self.ctx.read(d_buf, np.float16, (m, cfg.d_model)).astype(F32) - moe[tok] += d_ * down_s[e] * wts[tok, slot][:, None] + dn = self.ctx.read(d_buf, np.float16, (m_total, d)).astype(F32) + moe = np.zeros((N, d), dtype=F32) + for e, tok, slot, r0 in experts: + moe[tok] += dn[r0:r0 + len(tok)] * down_s[e] * wts[tok, slot][:, None] return moe # -- forward --------------------------------------------------------------- def forward(self, ids: np.ndarray, P: int, sc_logits: np.ndarray | None = None, - sc_temp_inv: float = 1.0, sc_use: float = 1.0) -> np.ndarray: - """sc_logits=None -> the Stage-1-validated zero-SC unified forward.""" + sc_temp_inv: float = 1.0, sc_use: float = 1.0, + sc_probs16: np.ndarray | None = None) -> np.ndarray: + """sc_logits=None and sc_probs16=None -> the Stage-1-validated zero-SC + unified forward.""" cfg, w = self.cfg, self.w ids = np.asarray(ids) N = len(ids) @@ -411,9 +441,10 @@ def forward(self, ids: np.ndarray, P: int, sc_logits: np.ndarray | None = None, dmp = self.dump or (lambda name, il, arr: None) sc_sig = None - if sc_logits is not None: + if sc_logits is not None or sc_probs16 is not None: with objc.autorelease_pool(): - sc_sig = self._sc_signal(sc_logits, sc_temp_inv, sc_use) + sc_sig = self._sc_signal(sc_logits, sc_temp_inv, sc_use, + probs16=sc_probs16) dmp("sc_sig", -1, sc_sig) x = embed_tokens(w, cfg, ids, P, sc_sig=sc_sig) dmp("inp_region", -1, x) @@ -480,6 +511,9 @@ def forward(self, ids: np.ndarray, P: int, sc_logits: np.ndarray | None = None, with objc.autorelease_pool(): logits = self._gemm_f32("token_embd.weight", x[P:], cfg.vocab_size) cap = F32(cfg.final_logit_softcap) - logits = (np.tanh(logits / cap) * cap).astype(F32) + # in-place softcap: the expression form held ~3 extra 268 MB transients + np.divide(logits, cap, out=logits) + np.tanh(logits, out=logits) + np.multiply(logits, cap, out=logits) dmp("result_output", -1, logits) return logits diff --git a/SuperKittens/models/gemma/diffusion/generate.py b/SuperKittens/models/gemma/diffusion/generate.py index b839fd4..5326266 100644 --- a/SuperKittens/models/gemma/diffusion/generate.py +++ b/SuperKittens/models/gemma/diffusion/generate.py @@ -70,24 +70,40 @@ def trim_canvas(canvas: np.ndarray, eog: set[int]) -> int: return cut +def probs16_of(logits: np.ndarray, temp_inv: float, _chunk: int = 32) -> np.ndarray: + """softmax(logits * temp_inv) -> f16, chunked; same math _sc_signal runs, + pulled forward so the raw 268 MB f32 block can be freed between forwards + (anonymous footprint drives swapfile growth on the 16 GB / tight-disk host).""" + from .graph_ref import softmax + C, V = logits.shape + out = np.empty((C, V), np.float16) + for c0 in range(0, C, _chunk): + c1 = min(c0 + _chunk, C) + out[c0:c1] = softmax(logits[c0:c1] * F32(temp_inv)).astype(np.float16) + return out + + def run_block(model, cfg, prompt_ids: np.ndarray, params: EBParams, use_sc: bool, log, mode: str = "gpu"): """One denoising block; returns (argmax_canvas, steps_run, timings).""" C = cfg.canvas_length P = len(prompt_ids) smp = EntropyBoundSampler(params, cfg.vocab_size, C) - prev_logits = None + prev_logits = None # cpu mode (oracle keeps the sc_logits contract) + prev_probs = None # gpu mode (precomputed softmax, raw logits freed) fw_s = smp_s = 0.0 step = None while not smp.finished: ids = np.concatenate([prompt_ids, smp.canvas.astype(np.int32)]) t0 = time.time() - if use_sc and prev_logits is not None: - # step 0 (sc_use=0 in the reference) is bit-identical to zero-SC: - # sig*0 == 0 and x+0 == x — skip the 1.5 GB embT stream entirely + if use_sc and prev_probs is not None: + logits = model.forward(ids, P, sc_probs16=prev_probs, sc_use=1.0) + elif use_sc and prev_logits is not None: logits = model.forward(ids, P, sc_logits=prev_logits, sc_temp_inv=float(smp.prev_temp_inv), sc_use=1.0) else: + # step 0 (sc_use=0 in the reference) is bit-identical to zero-SC: + # sig*0 == 0 and x+0 == x — skip the 1.5 GB embT stream entirely logits = model.forward(ids, P) t1 = time.time() if not np.isfinite(logits).all(): @@ -96,7 +112,12 @@ def run_block(model, cfg, prompt_ids: np.ndarray, params: EBParams, t2 = time.time() fw_s += t1 - t0 smp_s += t2 - t1 - prev_logits = logits + if use_sc and not smp.finished: + if mode == "gpu": + prev_probs = probs16_of(logits, float(smp.prev_temp_inv)) + else: + prev_logits = logits + del logits rec = {"step": step.step_idx, "t": round(step.t, 4), "accepted": int(step.accepted.sum()), "held": step.held, "H_mean": round(step.entropy_mean, 5), "finished": step.finished, From 42073383ad6a721c523b5c14ca3a59a8008167ec Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 01:10:04 -0400 Subject: [PATCH 25/31] preserve: diffgemma stage2 gate-3 prep WIP (agent stall recovery; gates 1-2 GREEN, e2e run pending) --- temp/diffgemma_s2/STATUS.md | 23 +++++++++++++++++++++++ temp/diffgemma_s2/run_gate3.sh | 2 ++ 2 files changed, 25 insertions(+) diff --git a/temp/diffgemma_s2/STATUS.md b/temp/diffgemma_s2/STATUS.md index 41cb09a..7d1df81 100644 --- a/temp/diffgemma_s2/STATUS.md +++ b/temp/diffgemma_s2/STATUS.md @@ -127,6 +127,29 @@ GATE5_PLACEHOLDER - `compare_eb.py`, `eb_ref_harness.cpp`, `rng_dump.cpp` — Stage-2a synthetic sampler gate (committed earlier, still pass). +## Memory war story, Stage-2 edition (multi-forward generation) + +Stage-1 stabilized ONE forward per process; generation runs many. First two +e2e attempts died during step 3 (first: external SIGKILL — kernel got there +first; second: the disk watchdog, correctly): system swap grew ~800 MB/30 s +during forwards, macOS swapfiles consumed the ~2 GB-free root volume, and the +box headed for the Stage-1 disk-exhaustion failure mode. Telemetry +(gen/*/mem.log): swap 1.57→2.35 GB in 30 s, disk free 2.1 GB→264 MB in 60 s. + +Fix (commit bcb5b61), forward verified bit-identical pre/post on both the +zero-SC and SC legs: +- persistent grow-only activation scratch in MetalCtx (per-call-site tag): + every _gemm_f32 A/C, attention q/k/vt/mask/s/p/o, and a packed-MoE rewrite + (all hit experts share three row-packed scratch buffers + offsets instead + of ~3 GB/forward of per-expert MTLBuffer alloc/free churn); +- generation loop converts each step's logits to the next step's SC probs + (f16) immediately and frees the raw 268 MB f32 block (`probs16_of`); +- in-place final softcap (the expression form held ~3 extra 268 MB + transients at the highest-pressure moment). +Also freed: Gate-1 logit artifacts (~1.3 GB) after the bit-regression diff. +Scratch reuse is also slightly faster: 43.8 s zero-SC / 50.8 s SC (was +47.5/52.2). + ## Host notes - amelia root volume runs ~2.5-4.8 GB free with the embT + dumps in place; diff --git a/temp/diffgemma_s2/run_gate3.sh b/temp/diffgemma_s2/run_gate3.sh index 170369e..94bf626 100644 --- a/temp/diffgemma_s2/run_gate3.sh +++ b/temp/diffgemma_s2/run_gate3.sh @@ -18,6 +18,8 @@ while kill -0 $PID 2>/dev/null; do sleep 30; SECS=$((SECS+30)) swap=$(sysctl -n vm.swapusage | awk '{print $6}' | tr -d M) free=$(df -m /System/Volumes/Data | tail -1 | awk '{print $4}') + rss=$(ps -o rss= -p $PID 2>/dev/null) + echo "t=$SECS rss_kb=$rss swap=${swap}M diskfree=${free}M freepct=$(memory_pressure -Q 2>/dev/null | awk -F': ' '/percentage/{print $2}')" >> $out/mem.log if (( ${swap%.*} > 5500 )); then echo "WATCHDOG swap=${swap}M KILL" >> $out/run.log; kill -9 $PID; exit 3; fi if (( free < 400 )); then echo "WATCHDOG diskfree=${free}M KILL" >> $out/run.log; kill -9 $PID; exit 4; fi if (( SECS > 3600 )); then echo "WATCHDOG timeout KILL" >> $out/run.log; kill -9 $PID; exit 5; fi From 37fd38eabc752b402497424c9fd8beadf25931da Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 01:10:21 -0400 Subject: [PATCH 26/31] preserve: granite launcher WIP (agent stall recovery) --- SuperKittens/models/granite/launcher.c++ | 6 +- build.log | 174 +++++++++++++++++++++++ 2 files changed, 178 insertions(+), 2 deletions(-) create mode 100644 build.log diff --git a/SuperKittens/models/granite/launcher.c++ b/SuperKittens/models/granite/launcher.c++ index d1b8fd3..6e78135 100644 --- a/SuperKittens/models/granite/launcher.c++ +++ b/SuperKittens/models/granite/launcher.c++ @@ -248,9 +248,11 @@ extern "C" int sk_granite_load_gguf(sk_granite_handle* hp, const char* path) { return -10; } - // Layer types: head_count_kv is a per-layer U32 array; 0 = mamba (recurrent). + // Layer types: head_count_kv is a per-layer int array (I32 in the official + // GGUFs); 0 = mamba (recurrent), >0 = attention. const auto* kvh = sk::gguf::meta_find(gm, "granitehybrid.attention.head_count_kv"); - if (!kvh || kvh->type != sk::gguf::V_ARRAY || kvh->arr_type != sk::gguf::V_U32 + if (!kvh || kvh->type != sk::gguf::V_ARRAY + || (kvh->arr_type != sk::gguf::V_U32 && kvh->arr_type != sk::gguf::V_I32) || kvh->arr_len != c.n_layers) { std::fprintf(stderr, "granite gguf: head_count_kv array missing/mismatched\n"); return -11; diff --git a/build.log b/build.log new file mode 100644 index 0000000..b57546a --- /dev/null +++ b/build.log @@ -0,0 +1,174 @@ +=== compiling Metal kernels === + kernels/fusion/gemv_geglu_bf16_m1.metal + kernels/fusion/bias_add.metal + kernels/fusion/rms_residual.metal + kernels/fusion/add_rmsnorm.metal + kernels/fusion/silu_mul.metal + kernels/fusion/gemv_bf16_m1.metal + kernels/fusion/q8_0_swiglu_prenorm_m1.metal + kernels/fusion/gated_mlp_gelu.metal + kernels/fusion/gemm_res_norm.metal + kernels/fusion/gemv_swiglu_m1.metal + kernels/fusion/gated_mlp.metal + kernels/fusion/add_rmsnorm_bf16.metal + kernels/fusion/kv_up_pair.metal + kernels/fusion/gated_mlp_bf16.metal + kernels/fusion/rms_rope.metal + kernels/fusion/gemm_bias_act.metal + kernels/utils/layernorm/layernorm.metal + kernels/utils/embedding/embedding.metal + kernels/utils/embedding/embedding_lookup_bf16.metal + kernels/utils/rmsnorm/rmsnorm_bf16.metal + kernels/utils/rmsnorm/rmsnorm_t1.metal + kernels/conv/conv1d.metal + kernels/conv/conv3d.metal + kernels/conv/conv2d.metal + kernels/gemm/q8_0_matvec_addres.metal + kernels/gemm/q2k_matvec.metal + kernels/gemm/q4k_matvec_bf16.metal + kernels/gemm/q6k_matvec.metal + kernels/gemm/q8_0_matvec.metal + kernels/gemm/gemm_mma.metal + kernels/gemm/q5k_matvec.metal + kernels/gemm/q8_0_swiglu_m1.metal + kernels/gemm/q8_0_matvec_bf16.metal + kernels/gemm/bf16/gemm.metal + kernels/gemm/q6k_matvec_bf16.metal + kernels/gemm/q4k_matvec.metal + kernels/gemm/iq2xxs_matvec.metal + kernels/gemm/fp16/gemv_t_2dtile.metal + kernels/gemm/fp16/gemv.metal + kernels/gemm/fp16/gemm.metal + kernels/gemm/fp16/gemv_t.metal + kernels/gemm/gemm_mma_smallm.metal + kernels/gemm/q3k_matvec.metal + kernels/gemm/fp8/gemm.metal + kernels/attn/attn.metal + kernels/rotary/rotary.metal + kernels/rotary/rope_qk_bf16.metal + kernels/ops/split/split_packed.metal + kernels/ops/cast/cast.metal + kernels/ops/transpose/seq_head.metal + kernels/ops/causal_mask/causal_mask.metal + kernels/ops/add/add.metal + kernels/ops/kv_cache/kv_cache.metal + kernels/ops/kv_cache/kv_cache_write_bf16.metal + kernels/ops/sample/argmax_2pass.metal + kernels/ops/sample/sample.metal + kernels/ops/sample/argmax_bf16.metal + kernels/moe/router.metal + kernels/moe/swiglu_pair_iq2xxs.metal + kernels/moe/down_scatter_q2k.metal + kernels/moe/router_v3.metal + kernels/moe/down_scatter.metal + kernels/moe/swiglu_pair.metal + kernels/flash_attn/flash_attn.metal + models/ssm/mamba2/conv1d_silu.metal + models/ssm/mamba2/mamba2_step.metal + models/ssm/mamba2/gate_norm.metal + models/ssm/mamba2/mamba2_step_ref.metal + models/ssm/mamba2/mamba2_ssd_ref.metal + models/ssm/mamba2/gevm_splitk.metal + models/ssm/mamba2/mamba2_ssd_chunked.metal + models/ssm/mamba2/mamba2_ssd.metal + models/ssm/mamba3/mamba3_step_post.metal + models/ssm/mamba3/pre_ssm.metal + models/ssm/mamba3/mamba3_ssm.metal + models/ssm/mamba3/post_ssm.metal + models/ssm/mamba3/mamba3_decode_full.metal + models/ssm/mamba3/mamba3_step.metal + models/gemma/gemma4/qkv_norm.metal + models/gemma/gemma4/ple_inject.metal + models/gemma/gemma4/gemma4_qkv_norm_rope_partial_t1.metal + models/gemma/gemma4/ple_inject_fused.metal + models/gemma/gemma4/gemma4_ops.metal + models/gemma/gemma4/prope.metal + models/gemma/gemma4/attn.metal + models/gemma/gemma4/gemma4_gemm_mma.metal + models/gemma/diffusion/dg_kernels.metal + ⚠ compile failed — skipping + models/granite/granite_ops.metal + models/qwen/qwen_rope.metal + models/deepseek/kernels/mla_v2.metal + models/deepseek/kernels/dsv4_kv.metal + models/deepseek/kernels/repeat.metal + models/deepseek/kernels/sum_rows.metal + models/deepseek/kernels/dsv4_hc.metal + models/deepseek/kernels/moe_group.metal + models/deepseek/kernels/set_rows.metal + models/deepseek/kernels/mla_glue.metal + models/deepseek/kernels/moe_group_mma.metal + models/deepseek/kernels/rope_interleave.metal + models/deepseek/kernels/moe_glue.metal + models/deepseek/kernels/concat.metal + models/deepseek/kernels/shared_swiglu.metal + models/deepseek/kernels/dsv4_rope.metal + models/deepseek/kernels/dsv4_misc.metal + models/deepseek/kernels/cpy.metal + models/deepseek/kernels/mul_mv_id.metal + models/deepseek/kernels/softmax.metal + models/deepseek/kernels/argsort.metal + models/deepseek/kernels/bin.metal +=== linking metallib === + → build/libsk.metallib +=== compiling C dispatchers === + SuperKittens/kernels/metal_impl.cpp (metal-cpp impls) +SuperKittens/kernels/metal_impl.cpp:3:9: warning: 'NS_PRIVATE_IMPLEMENTATION' macro redefined [-Wmacro-redefined] + 3 | #define NS_PRIVATE_IMPLEMENTATION + | ^ +:1:9: note: previous definition is here + 1 | #define NS_PRIVATE_IMPLEMENTATION 1 + | ^ +SuperKittens/kernels/metal_impl.cpp:4:9: warning: 'MTL_PRIVATE_IMPLEMENTATION' macro redefined [-Wmacro-redefined] + 4 | #define MTL_PRIVATE_IMPLEMENTATION + | ^ +:2:9: note: previous definition is here + 2 | #define MTL_PRIVATE_IMPLEMENTATION 1 + | ^ +SuperKittens/kernels/metal_impl.cpp:5:9: warning: 'CA_PRIVATE_IMPLEMENTATION' macro redefined [-Wmacro-redefined] + 5 | #define CA_PRIVATE_IMPLEMENTATION + | ^ +:3:9: note: previous definition is here + 3 | #define CA_PRIVATE_IMPLEMENTATION 1 + | ^ +3 warnings generated. + SuperKittens/kernels/fusion/kv_up_pair.c++ + SuperKittens/kernels/fusion/add_rmsnorm.c++ + SuperKittens/kernels/utils/layernorm/layernorm.c++ + SuperKittens/kernels/utils/embedding/embedding.c++ + SuperKittens/kernels/conv/conv.c++ + SuperKittens/kernels/gemm/quant_mv.c++ + SuperKittens/kernels/gemm/gemm.c++ + SuperKittens/kernels/attn/attn.c++ + SuperKittens/kernels/rotary/rope_tail.c++ + SuperKittens/kernels/rotary/rotary.c++ + SuperKittens/kernels/moe/swiglu_pair.c++ + SuperKittens/kernels/moe/down_scatter_q2k.c++ + SuperKittens/kernels/moe/down_scatter.c++ + SuperKittens/kernels/moe/router.c++ + SuperKittens/kernels/moe/swiglu_pair_iq2xxs.c++ + SuperKittens/kernels/flash_attn/flash_attn.c++ + SuperKittens/models/ssm/mamba2/launcher.c++ + SuperKittens/models/ssm/mamba2/mamba2.c++ + SuperKittens/models/ssm/mamba2/weights.c++ + SuperKittens/models/ssm/mamba3/mamba3.c++ + SuperKittens/models/gemma/gemma4/launcher.c++ + SuperKittens/models/gemma/gemma4/weights.c++ + SuperKittens/models/granite/launcher.c++ + SuperKittens/models/qwen/launcher.c++ + SuperKittens/models/qwen/weights.c++ + SuperKittens/models/deepseek/launcher.c++ + SuperKittens/models/deepseek/weights.c++ + SuperKittens/models/load/gguf/gguf.c++ + SuperKittens/models/load/safetensor/safetensor.c++ + SuperKittens/inference/weight_store.c++ + SuperKittens/inference/quantize.c++ + SuperKittens/inference/quantize_q4k.c++ + SuperKittens/inference/silicon/mmap_buffer.c++ + SuperKittens/inference/silicon/icb_recorder.c++ +=== linking dylib === + → build/libsk.dylib + +Build complete. Files in build/: +-rwxr-xr-x@ 1 alazarmanakelew staff 791K Jun 11 00:25 build/libsk.dylib +-rw-r--r--@ 1 alazarmanakelew staff 1.3M Jun 11 00:25 build/libsk.metallib From c5321c54d2ad496d1a20ff6e33be7331a9e9c58e Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 01:16:43 -0400 Subject: [PATCH 27/31] diffgemma stage2 gate-3 runners: tighten watchdog to swap>3.5G/disk<4G/90s polls; watchdog the gate-4 cli leg --- temp/diffgemma_s2/run_gate3.sh | 11 ++++++----- temp/diffgemma_s2/run_queue2.sh | 28 ++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) create mode 100644 temp/diffgemma_s2/run_queue2.sh diff --git a/temp/diffgemma_s2/run_gate3.sh b/temp/diffgemma_s2/run_gate3.sh index 94bf626..2f2b81e 100644 --- a/temp/diffgemma_s2/run_gate3.sh +++ b/temp/diffgemma_s2/run_gate3.sh @@ -1,7 +1,8 @@ #!/bin/zsh # run_gate3.sh [extra generate.py args...] -# One SK e2e generation, detached-safe, with the Stage-1 watchdog rules: -# kill on system swap > 5.5 GB, root-disk free < 400 MB, or wall > 3600 s. +# One SK e2e generation, detached-safe. Watchdog: system swap > 3.5 GB, +# root-disk free < 4 GB, or wall > 3600 s (post-cleanup amelia has ~12 GB +# free; tighter-than-Stage-1 thresholds abort long before the host is at risk). LAB=~/sk-diffg-s2b GGUF=~/diffgemma-gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf name=$1; pids=$2; steps=$3; seed=$4; shift 4 @@ -15,13 +16,13 @@ PID=$! echo $PID > $out/pid SECS=0 while kill -0 $PID 2>/dev/null; do - sleep 30; SECS=$((SECS+30)) + sleep 90; SECS=$((SECS+90)) swap=$(sysctl -n vm.swapusage | awk '{print $6}' | tr -d M) free=$(df -m /System/Volumes/Data | tail -1 | awk '{print $4}') rss=$(ps -o rss= -p $PID 2>/dev/null) echo "t=$SECS rss_kb=$rss swap=${swap}M diskfree=${free}M freepct=$(memory_pressure -Q 2>/dev/null | awk -F': ' '/percentage/{print $2}')" >> $out/mem.log - if (( ${swap%.*} > 5500 )); then echo "WATCHDOG swap=${swap}M KILL" >> $out/run.log; kill -9 $PID; exit 3; fi - if (( free < 400 )); then echo "WATCHDOG diskfree=${free}M KILL" >> $out/run.log; kill -9 $PID; exit 4; fi + if (( ${swap%.*} > 3500 )); then echo "WATCHDOG swap=${swap}M KILL" >> $out/run.log; kill -9 $PID; exit 3; fi + if (( free < 4000 )); then echo "WATCHDOG diskfree=${free}M KILL" >> $out/run.log; kill -9 $PID; exit 4; fi if (( SECS > 3600 )); then echo "WATCHDOG timeout KILL" >> $out/run.log; kill -9 $PID; exit 5; fi done wait $PID diff --git a/temp/diffgemma_s2/run_queue2.sh b/temp/diffgemma_s2/run_queue2.sh new file mode 100644 index 0000000..a131fb9 --- /dev/null +++ b/temp/diffgemma_s2/run_queue2.sh @@ -0,0 +1,28 @@ +#!/bin/zsh +# run_queue2.sh — Gate-3 remaining prompts + Gate-4 reference cli, strictly +# sequential (one heavy process at a time on the 16 GB host). +LAB=~/sk-diffg-s2b +GGUF=~/diffgemma-gguf/diffusiongemma-26B-A4B-it-Q4_K_M.gguf + +$LAB/tools/run_gate3.sh pq_s16 $LAB/gen/pqcli_prompt.i32 16 1234 +$LAB/tools/run_gate3.sh p2_s16 $LAB/gen/p2cli_prompt.i32 16 1234 + +# Gate 4 reference: uninstrumented cli (no DG_EB_DUMP in env), same prompt, +# same seed and S as the SK p1 run. Same watchdog rules as run_gate3.sh — +# the CPU cli is ~30 s/step so 3600 s is generous. +caffeinate -is ~/llamacpp-diffg/build/bin/llama-diffusion-cli -m $GGUF \ + -p "What is the capital of France?" \ + --diffusion-eb-max-steps 16 --seed 1234 -n 256 \ + > $LAB/gen/cli_p1_s16.log 2>&1 & +CPID=$! +SECS=0 +while kill -0 $CPID 2>/dev/null; do + sleep 90; SECS=$((SECS+90)) + swap=$(sysctl -n vm.swapusage | awk '{print $6}' | tr -d M) + free=$(df -m /System/Volumes/Data | tail -1 | awk '{print $4}') + echo "t=$SECS swap=${swap}M diskfree=${free}M" >> $LAB/gen/cli_p1_s16.mem.log + if (( ${swap%.*} > 3500 )); then echo "WATCHDOG swap KILL" >> $LAB/gen/cli_p1_s16.log; kill -9 $CPID; break; fi + if (( free < 4000 )); then echo "WATCHDOG disk KILL" >> $LAB/gen/cli_p1_s16.log; kill -9 $CPID; break; fi + if (( SECS > 3600 )); then echo "WATCHDOG timeout KILL" >> $LAB/gen/cli_p1_s16.log; kill -9 $CPID; break; fi +done +echo QUEUE2_DONE From ad77999a61d5d137aaafec007a17859056452f73 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 01:18:30 -0400 Subject: [PATCH 28/31] diffgemma stage2 GATE 3 p1 GREEN: e2e SK loop coherent on 'capital of France' (12 steps adaptive stop, 0 mask tokens, finite logits, correct answer; swap flat ~1.1G over 12 forwards) --- temp/diffgemma_s2/STATUS.md | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/temp/diffgemma_s2/STATUS.md b/temp/diffgemma_s2/STATUS.md index 7d1df81..60748d2 100644 --- a/temp/diffgemma_s2/STATUS.md +++ b/temp/diffgemma_s2/STATUS.md @@ -102,7 +102,29 @@ Reference cli's own final answer (same run): thought channel reasoning + ## Gate 3 — e2e coherent generation -GATE3_PLACEHOLDER +Full SK loop (Metal forward + SC + sampler) via generate.py under +tools/run_gate3.sh (caffeinate -is, detached, watchdog swap>3.5G / +disk<4G / 3600s). All prompts use the reference cli's own chat template +(`<|turn>system\n<|think|>\n\n<|turn>user\n{msg}\n +<|turn>model\n`); every id file was detokenized against +tokenizer.ggml.tokens and matched the intended text exactly. S=16, +seed=1234, C=256, SC on. Lab tree verified == repo tip by md5 before runs. + +### p1 "What is the capital of France?" (P=23) — run gen/p1_s16 + +12 steps (adaptive stop), wall 604.8 s, fw 593.8 s, sampler 7.7 s, +0 mask tokens in canvas, logits finite every step (generate.py raises +otherwise), trim 45. H_mean 0.464 -> 0.0017, accepted 126 -> 252, +swap flat ~1.1 G, disk drift ~170 MB over the run — the bcb5b61 scratch +fix holds over a full 12-forward generation. Output (verbatim): + +> `<|channel>thought` +> `The user is asking for the capital of France.` +> ` * Identify country: France.` +> ` * Identify city: Paris.` +> `State the answer clearly.The capital of France is Paris.` + +GATE3_MORE_PLACEHOLDER ## Gate 4 — cross-check vs llama-diffusion-cli From 75f1ad6022936d02535141a82e101a5da76d7b59 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 06:58:28 -0400 Subject: [PATCH 29/31] =?UTF-8?q?diffgemma=20stage2=20GATES=203+4=20GREEN?= =?UTF-8?q?=20+=20stage-3=20baseline:=20pq=20(2+2=3D4,=208=20steps)=20+=20?= =?UTF-8?q?p2=20(full-canvas=20haiku=20w/=20self-critique;=20trim=20heuris?= =?UTF-8?q?tic=20false-cuts=20comma=20lists)=20coherent;=20cli=20cross-che?= =?UTF-8?q?ck=20agrees=20on=20p1=20(Paris,=208=20steps);=20baseline=2049.5?= =?UTF-8?q?-50.4=20s/step=20wall,=20fw=2098%=20=E2=80=94=20CPU=20cli=2029.?= =?UTF-8?q?85=20s/step=20beats=20SK=20GPU=201.65x,=20Stage-3=20surface=20i?= =?UTF-8?q?s=20the=20forward?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- temp/diffgemma_s2/STATUS.md | 103 ++++++++++++++++++++++++++++++++---- 1 file changed, 94 insertions(+), 9 deletions(-) diff --git a/temp/diffgemma_s2/STATUS.md b/temp/diffgemma_s2/STATUS.md index 60748d2..3c07f0f 100644 --- a/temp/diffgemma_s2/STATUS.md +++ b/temp/diffgemma_s2/STATUS.md @@ -100,7 +100,7 @@ possible in principle on exact f32 ties; the synthetic gate quantified it at Reference cli's own final answer (same run): thought channel reasoning + "The capital of France is Paris." — 7 steps, 30.5 s/step CPU. -## Gate 3 — e2e coherent generation +## Gate 3 — e2e coherent generation: GREEN (3/3 prompts) Full SK loop (Metal forward + SC + sampler) via generate.py under tools/run_gate3.sh (caffeinate -is, detached, watchdog swap>3.5G / @@ -124,15 +124,94 @@ fix holds over a full 12-forward generation. Output (verbatim): > ` * Identify city: Paris.` > `State the answer clearly.The capital of France is Paris.` -GATE3_MORE_PLACEHOLDER +### pq "What is 2+2?" (P=23) — run gen/pq_s16 -## Gate 4 — cross-check vs llama-diffusion-cli +8 steps (adaptive stop), wall 395.6 s, fw 388.4 s, sampler 5.1 s, 0 mask +tokens, trim 48, EXIT 0. H_mean 0.576 -> 0.0015, accepted 98 -> 256. +Canvas beyond the cut is pure `` padding. Output (verbatim): -GATE4_PLACEHOLDER +> `<|channel>thought` +> `The user is asking for the sum of 2 and 2. This is a basic arithmetic question.` +> `2 + 2 = 4.` +> `Provide the answer directly.2 + 2 = 4` + +### p2 "Write a haiku about the ocean." (P=24) — run gen/p2_s16 + +14 steps (adaptive stop), wall 695.5 s, fw 682.4 s, sampler 9.2 s, 0 mask +tokens, EXIT 0. Creative prompt = high-entropy start: H_mean 2.75 -> 1.3e-4, +accepted 7 -> 256 (vs 126/98 step-0 accepts on the QA prompts). The model +filled the entire 256-token canvas (no EOS): brainstorm list, three haiku +drafts with self-critique and per-line syllable counts, closed the thought +channel, and began the final answer — truncated by canvas length, not by +incoherence. Full-canvas detok (verbatim, condensed): + +> `<|channel>thought` … `* Waves, tides, salt, blue, deep, vast, sand, +> shore, shells, crashing, rhythmic, endless, blue, foam.` … three drafted +> haikus each with `*Critique:*` … `* *Syllable count:*` … +> `Crashing waves on shore, / Endless secrets in the deep, / Salt mist fills +> the air.Crashing waves on shore,` + +`output.txt` shows only the first 31 tokens: `trim_canvas`'s stride-2 +repetition heuristic false-fires on the comma-separated brainstorm list +(the comma token repeats at stride 2 ≥ 6 times). Detok artifact in the +reference-cli-style trim, NOT a generation defect — keep in mind for +Stage 3 if trimmed output looks short on listy generations. + +**Gate 3 GREEN on all three prompts**: fluent prompt-relevant text, both +simple-QA prompts answered correctly, 0 mask tokens in every final canvas, +logits finite every step (generate.py raises otherwise), swap flat ~1.0-1.1 G +and disk free ~11.7 G throughout (bcb5b61 scratch fix holds across 8/12/14 +forward generations). + +## Gate 4 — cross-check vs llama-diffusion-cli: GREEN + +**GREEN — qualitative agreement.** Uninstrumented `llama-diffusion-cli` +(CPU, prompt-KV cache on), same prompt/seed/S as the SK p1 run +(`-p "What is the capital of France?" --diffusion-eb-max-steps 16 +--seed 1234 -n 256`), run gen/cli_p1_s16.log: + +> `<|channel>thought` +> `The user is asking for the capital of France.` +> ` * Entity: France.` +> ` * Question: Capital.` +> `The capital of France is Paris.The capital of France is Paris.` + +8 steps, 238.8 s total, 29.85 s/step. Same final answer and same +thought-channel structure as SK p1; trajectories diverge in step count +(8 vs 12) and thought wording as expected — the cli's chat template +prepends the system/think turns identically, but GPU f16 forward numerics +differ from CPU f32, so per-step accept sets drift after step 1 while both +converge to the correct answer. (Token-level sampler identity on shared +logits is already proven by Gate 2.) Swap flat 1044 M, disk ≥ 11.7 G +during the run. ## Stage-3 baseline numbers -GATE5_PLACEHOLDER +All runs: amelia M4 16 GB, C=256, S=16, seed 1234, SC on, Q4_K_M 26B-A4B. + +| run | prompt | steps used | wall (s) | wall/step | fw (s) | sampler (s) | fw share | +|---|---|---|---|---|---|---|---| +| p1_s16 | capital of France (P=23) | 12 | 604.8 | 50.4 | 593.8 | 7.7 | 98.2% | +| pq_s16 | 2+2 (P=23) | 8 | 395.6 | 49.5 | 388.4 | 5.1 | 98.2% | +| p2_s16 | ocean haiku (P=24) | 14 | 695.5 | 49.7 | 682.4 | 9.2 | 98.1% | + +- Per-step forward: step 0 (zero-SC) 44.9-46.0 s; steps ≥ 1 (SC active) + 48.5-52.2 s, median ~49 s — consistent with the 43.8/50.8 s bit-regression + timings, so queue-window contention cost ≲ a few percent. +- Host sampler 0.6-0.7 s/step; residual (probs16 conversion, IO, detok) + ~0.3 s/step. **The Metal forward is 98% of wall — it is the entire + Stage-3 surface.** Encode/sampler-side levers are noise at this split. +- Adaptive stop is doing real work: 8/12/14 steps used of S=16, and steps + scale with prompt entropy (QA converges fastest, creative slowest). +- Reference CPU cli on the same box: 29.85 s/step — the CPU reference + currently BEATS the SK GPU forward ~1.65× per step. The SK forward + (T=279 prefill-shaped, per-step full-canvas recompute, MoE expert + streaming over a 15 GB-resident Q4_K_M model) is unoptimized + Stage-2 correctness plumbing; closing (then inverting) that gap is the + Stage-3 objective. Candidate levers, in lab-evidence order: prompt-KV + reuse across steps (the cli already does this), Q4_K MoE GEMM port + (~4× over fp16 swiglu_pair, proven out-of-tree), gemm_mma for the + T=279 dense projections. ## Tools (this dir, all run on amelia from ~/sk-diffg-s2b) @@ -145,7 +224,8 @@ GATE5_PLACEHOLDER through the SK sampler, diffs every decision field, deletes consumed logits. - `make_embt.py` — one-time [d_model, vocab] f16 transposed embed (1.48 GB, amelia ~/sk-diffg-s2b/dg_embT_f16.bin) for the GPU SC soft-embed stream. -- `run_gate3.sh` — watchdogged e2e generation (swap>5.5G / disk<400M / 3600s). +- `run_gate3.sh` — watchdogged e2e generation (swap>3.5G / disk<4G / 3600s). +- `run_queue2.sh` — sequential pq/p2 gate-3 runs + watchdogged gate-4 cli. - `compare_eb.py`, `eb_ref_harness.cpp`, `rng_dump.cpp` — Stage-2a synthetic sampler gate (committed earlier, still pass). @@ -176,8 +256,13 @@ Scratch reuse is also slightly faster: 43.8 s zero-SC / 50.8 s SC (was - amelia root volume runs ~2.5-4.8 GB free with the embT + dumps in place; every logits file is 268 MB — delete as consumed (gate2 streams + deletes). -- A GGUF→derek transfer (`cat ~/diffgemma-gguf/...gguf`) was running through - amelia during the correctness runs; timing-sensitive numbers (Gate 5) were - taken TRANSFER_NOTE_PLACEHOLDER. +- A GGUF→derek transfer (chunked `tail -c +N ~/diffgemma-gguf/...gguf`) + was streaming from amelia around the Gate-3 window (chunk resumed at byte + 4.2e9 at 01:09:35, definitively concurrent with the pq/p2/cli runs; p1 + overlap uncertain). It consumes network + disk-read only; per-step fw_s is + consistent across all three SK runs (and with the pre-queue 43.8/50.8 s + bit-regression timings), so the Gate-5 numbers carry at most a few percent + of contention noise. A ~3.6 GB-RSS Virtualization VM was also resident the + whole time. - Stale `/Users/amelia/SuperKittens` partial copy exists; the lab runs pin PYTHONPATH=~/sk-diffg-s2b so the rsynced tree wins. Don't import without it. From 551ab1e41aca06f2e3ee8923d0caa38c3d69133d Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 07:19:06 -0400 Subject: [PATCH 30/31] =?UTF-8?q?granite:=20Stage-1=20gates=20green=20?= =?UTF-8?q?=E2=80=94=20e2e=20COHERENT=20on=20lexie=20(QA=2048/48=20token-G?= =?UTF-8?q?OLD=20vs=20llama.cpp=20CPU=20greedy;=20poem=20fp16=20near-tie?= =?UTF-8?q?=20divergence=20@8,=20gap=200.125;=20qwen3-1.7b=20no-regression?= =?UTF-8?q?=2032/32=20vs=20pristine=20main;=2051.2=20tok/s=20median=20deco?= =?UTF-8?q?de)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SuperKittens/models/granite/STATUS.md | 34 +++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 SuperKittens/models/granite/STATUS.md diff --git a/SuperKittens/models/granite/STATUS.md b/SuperKittens/models/granite/STATUS.md new file mode 100644 index 0000000..23c51cf --- /dev/null +++ b/SuperKittens/models/granite/STATUS.md @@ -0,0 +1,34 @@ +# Granite-4 hybrid — Stage 1: PORTED-COHERENT (2026-06-11) + +Model: **granite-4.0-h-1b** (Q8_0, arch `granitehybrid`) — 40 layers: 36 mamba2 ++ 4 attention (A at 5/15/25/35, per-layer type from +`granitehybrid.attention.head_count_kv`), dense SwiGLU FFN every layer, NoPE +attention (attention_multiplier 1/128), granite scalar multipliers +(embed 12 / residual 0.22 / logit 1/6), tied Q8_0 head, vocab 100352. +h-1b over h-micro because h-micro is head_dim=64 (SK dense attention is D=128-only). + +Reuses mamba2 family kernels (conv1d_silu/_step, conv_state_capture, +mamba2_ssd, gate_norm) and shared dense kernels (q8_0_matvec, mha_causal D=128, +kv_cache_write, rmsnorm, silu_mul) unchanged; `granite_ops.metal` adds the two +elementwise kernels that carry the granite multipliers. + +## Gates (lexie M4 base, greedy) +- **Load**: every dim cross-checked against GGUF metadata, fail-loud; config + table printed at load. +- **Coherence**: QA prompt is **48/48 token-identical to llama.cpp CPU greedy** + (GGML_METAL=OFF, GGML_CPU_REPACK=OFF). Pizza-dough poem diverges at token 8 + on an fp16 near-tie (SK top-2 gap 0.125: ' cr' 36.4062 vs ' mound' 36.2812); + both continuations fluent. Logits finite. Outputs byte-identical across + M-series hosts. +- **No-regression**: Qwen3-1.7B-Q8 32-tok greedy token-identical (32/32), + branch build vs pristine main build. +- **Decode**: median **51.2 tok/s** (2 warmup + 5 reps × 64 tok; spread + 50.99–51.30; indicative — lexie is thermal-drifty). + +## Stage-1 limits +- Prompt must fit one prefill forward (chunked mamba prefill would re-zero-pad + the conv left edge); decode is O(1)-state. +- `generate_n` demands a fresh sequence (mamba state); the adapter resets per + call. +- fp16 logits end-to-end: greedy near-ties (gap ≲ 0.13) may flip vs f32 + references. From f3c688b9e902b0685440b73a61bd4a163b68d7f4 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Thu, 11 Jun 2026 07:20:14 -0400 Subject: [PATCH 31/31] granite: drop stray build.log from repo root --- build.log | 174 ------------------------------------------------------ 1 file changed, 174 deletions(-) delete mode 100644 build.log diff --git a/build.log b/build.log deleted file mode 100644 index b57546a..0000000 --- a/build.log +++ /dev/null @@ -1,174 +0,0 @@ -=== compiling Metal kernels === - kernels/fusion/gemv_geglu_bf16_m1.metal - kernels/fusion/bias_add.metal - kernels/fusion/rms_residual.metal - kernels/fusion/add_rmsnorm.metal - kernels/fusion/silu_mul.metal - kernels/fusion/gemv_bf16_m1.metal - kernels/fusion/q8_0_swiglu_prenorm_m1.metal - kernels/fusion/gated_mlp_gelu.metal - kernels/fusion/gemm_res_norm.metal - kernels/fusion/gemv_swiglu_m1.metal - kernels/fusion/gated_mlp.metal - kernels/fusion/add_rmsnorm_bf16.metal - kernels/fusion/kv_up_pair.metal - kernels/fusion/gated_mlp_bf16.metal - kernels/fusion/rms_rope.metal - kernels/fusion/gemm_bias_act.metal - kernels/utils/layernorm/layernorm.metal - kernels/utils/embedding/embedding.metal - kernels/utils/embedding/embedding_lookup_bf16.metal - kernels/utils/rmsnorm/rmsnorm_bf16.metal - kernels/utils/rmsnorm/rmsnorm_t1.metal - kernels/conv/conv1d.metal - kernels/conv/conv3d.metal - kernels/conv/conv2d.metal - kernels/gemm/q8_0_matvec_addres.metal - kernels/gemm/q2k_matvec.metal - kernels/gemm/q4k_matvec_bf16.metal - kernels/gemm/q6k_matvec.metal - kernels/gemm/q8_0_matvec.metal - kernels/gemm/gemm_mma.metal - kernels/gemm/q5k_matvec.metal - kernels/gemm/q8_0_swiglu_m1.metal - kernels/gemm/q8_0_matvec_bf16.metal - kernels/gemm/bf16/gemm.metal - kernels/gemm/q6k_matvec_bf16.metal - kernels/gemm/q4k_matvec.metal - kernels/gemm/iq2xxs_matvec.metal - kernels/gemm/fp16/gemv_t_2dtile.metal - kernels/gemm/fp16/gemv.metal - kernels/gemm/fp16/gemm.metal - kernels/gemm/fp16/gemv_t.metal - kernels/gemm/gemm_mma_smallm.metal - kernels/gemm/q3k_matvec.metal - kernels/gemm/fp8/gemm.metal - kernels/attn/attn.metal - kernels/rotary/rotary.metal - kernels/rotary/rope_qk_bf16.metal - kernels/ops/split/split_packed.metal - kernels/ops/cast/cast.metal - kernels/ops/transpose/seq_head.metal - kernels/ops/causal_mask/causal_mask.metal - kernels/ops/add/add.metal - kernels/ops/kv_cache/kv_cache.metal - kernels/ops/kv_cache/kv_cache_write_bf16.metal - kernels/ops/sample/argmax_2pass.metal - kernels/ops/sample/sample.metal - kernels/ops/sample/argmax_bf16.metal - kernels/moe/router.metal - kernels/moe/swiglu_pair_iq2xxs.metal - kernels/moe/down_scatter_q2k.metal - kernels/moe/router_v3.metal - kernels/moe/down_scatter.metal - kernels/moe/swiglu_pair.metal - kernels/flash_attn/flash_attn.metal - models/ssm/mamba2/conv1d_silu.metal - models/ssm/mamba2/mamba2_step.metal - models/ssm/mamba2/gate_norm.metal - models/ssm/mamba2/mamba2_step_ref.metal - models/ssm/mamba2/mamba2_ssd_ref.metal - models/ssm/mamba2/gevm_splitk.metal - models/ssm/mamba2/mamba2_ssd_chunked.metal - models/ssm/mamba2/mamba2_ssd.metal - models/ssm/mamba3/mamba3_step_post.metal - models/ssm/mamba3/pre_ssm.metal - models/ssm/mamba3/mamba3_ssm.metal - models/ssm/mamba3/post_ssm.metal - models/ssm/mamba3/mamba3_decode_full.metal - models/ssm/mamba3/mamba3_step.metal - models/gemma/gemma4/qkv_norm.metal - models/gemma/gemma4/ple_inject.metal - models/gemma/gemma4/gemma4_qkv_norm_rope_partial_t1.metal - models/gemma/gemma4/ple_inject_fused.metal - models/gemma/gemma4/gemma4_ops.metal - models/gemma/gemma4/prope.metal - models/gemma/gemma4/attn.metal - models/gemma/gemma4/gemma4_gemm_mma.metal - models/gemma/diffusion/dg_kernels.metal - ⚠ compile failed — skipping - models/granite/granite_ops.metal - models/qwen/qwen_rope.metal - models/deepseek/kernels/mla_v2.metal - models/deepseek/kernels/dsv4_kv.metal - models/deepseek/kernels/repeat.metal - models/deepseek/kernels/sum_rows.metal - models/deepseek/kernels/dsv4_hc.metal - models/deepseek/kernels/moe_group.metal - models/deepseek/kernels/set_rows.metal - models/deepseek/kernels/mla_glue.metal - models/deepseek/kernels/moe_group_mma.metal - models/deepseek/kernels/rope_interleave.metal - models/deepseek/kernels/moe_glue.metal - models/deepseek/kernels/concat.metal - models/deepseek/kernels/shared_swiglu.metal - models/deepseek/kernels/dsv4_rope.metal - models/deepseek/kernels/dsv4_misc.metal - models/deepseek/kernels/cpy.metal - models/deepseek/kernels/mul_mv_id.metal - models/deepseek/kernels/softmax.metal - models/deepseek/kernels/argsort.metal - models/deepseek/kernels/bin.metal -=== linking metallib === - → build/libsk.metallib -=== compiling C dispatchers === - SuperKittens/kernels/metal_impl.cpp (metal-cpp impls) -SuperKittens/kernels/metal_impl.cpp:3:9: warning: 'NS_PRIVATE_IMPLEMENTATION' macro redefined [-Wmacro-redefined] - 3 | #define NS_PRIVATE_IMPLEMENTATION - | ^ -:1:9: note: previous definition is here - 1 | #define NS_PRIVATE_IMPLEMENTATION 1 - | ^ -SuperKittens/kernels/metal_impl.cpp:4:9: warning: 'MTL_PRIVATE_IMPLEMENTATION' macro redefined [-Wmacro-redefined] - 4 | #define MTL_PRIVATE_IMPLEMENTATION - | ^ -:2:9: note: previous definition is here - 2 | #define MTL_PRIVATE_IMPLEMENTATION 1 - | ^ -SuperKittens/kernels/metal_impl.cpp:5:9: warning: 'CA_PRIVATE_IMPLEMENTATION' macro redefined [-Wmacro-redefined] - 5 | #define CA_PRIVATE_IMPLEMENTATION - | ^ -:3:9: note: previous definition is here - 3 | #define CA_PRIVATE_IMPLEMENTATION 1 - | ^ -3 warnings generated. - SuperKittens/kernels/fusion/kv_up_pair.c++ - SuperKittens/kernels/fusion/add_rmsnorm.c++ - SuperKittens/kernels/utils/layernorm/layernorm.c++ - SuperKittens/kernels/utils/embedding/embedding.c++ - SuperKittens/kernels/conv/conv.c++ - SuperKittens/kernels/gemm/quant_mv.c++ - SuperKittens/kernels/gemm/gemm.c++ - SuperKittens/kernels/attn/attn.c++ - SuperKittens/kernels/rotary/rope_tail.c++ - SuperKittens/kernels/rotary/rotary.c++ - SuperKittens/kernels/moe/swiglu_pair.c++ - SuperKittens/kernels/moe/down_scatter_q2k.c++ - SuperKittens/kernels/moe/down_scatter.c++ - SuperKittens/kernels/moe/router.c++ - SuperKittens/kernels/moe/swiglu_pair_iq2xxs.c++ - SuperKittens/kernels/flash_attn/flash_attn.c++ - SuperKittens/models/ssm/mamba2/launcher.c++ - SuperKittens/models/ssm/mamba2/mamba2.c++ - SuperKittens/models/ssm/mamba2/weights.c++ - SuperKittens/models/ssm/mamba3/mamba3.c++ - SuperKittens/models/gemma/gemma4/launcher.c++ - SuperKittens/models/gemma/gemma4/weights.c++ - SuperKittens/models/granite/launcher.c++ - SuperKittens/models/qwen/launcher.c++ - SuperKittens/models/qwen/weights.c++ - SuperKittens/models/deepseek/launcher.c++ - SuperKittens/models/deepseek/weights.c++ - SuperKittens/models/load/gguf/gguf.c++ - SuperKittens/models/load/safetensor/safetensor.c++ - SuperKittens/inference/weight_store.c++ - SuperKittens/inference/quantize.c++ - SuperKittens/inference/quantize_q4k.c++ - SuperKittens/inference/silicon/mmap_buffer.c++ - SuperKittens/inference/silicon/icb_recorder.c++ -=== linking dylib === - → build/libsk.dylib - -Build complete. Files in build/: --rwxr-xr-x@ 1 alazarmanakelew staff 791K Jun 11 00:25 build/libsk.dylib --rw-r--r--@ 1 alazarmanakelew staff 1.3M Jun 11 00:25 build/libsk.metallib