models+gemm: Phi-4 and Granite-4-hybrid families, DiffusionGemma e2e (parity+coherent), gemma4 batched prefill (5.1-5.5x), fp16-GEMM row-grid fixes

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.
@@ -381,6 +402,41 @@ 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.
+    # 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",
+    ),
 }
 
 

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();

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"]

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)

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

SuperKittens/models/gemma/diffusion/dg_kernels.metal

@@ -0,0 +1,176 @@
+// dg_kernels.metal — DiffusionGemma family kernels.
+//
+// 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.
+//
+// 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.
+
+// ── 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 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]])
+{
+    const uint tid  = tid3.x;
+    const uint tptg = tptg3.x;
+    const uint r = tgpig.x;
+    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, 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(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);
+
+    for (uint c = tid; c < ncols; c += tptg) {
+        prow[c] = (half)(exp(row[c] * scale + mrow[c] - rmax) * inv);
+    }
+}