| """ClefRowPrefill: the row graph of this conversion = one state-free causal row of the Qwen3.5-9B text stack of |
| Cloudflare/clef-flash (revision f0b7d3cd), embeddings and position tables in, the final-norm hidden state of every |
| position out. |
| |
| forward(embeds f32 [1, L, 4096], valid f32 [1, L], cos f32 [1, L, 64], sin f32 [1, L, 64]) -> {"hidden": f32 [1, L, 4096]} |
| |
| embeds = the checkpoint's bf16 `embed_tokens` rows widened to float32 by the host (pad rows 0); valid = 1.0 on the |
| record's tokens, 0.0 on the right pads; cos / sin = the host's M-RoPE tables (scripts/clef_host.py; a text record = |
| positions (i, i, i)); hidden = last_hidden_state after the final RMSNorm at every position (what the author's joint |
| schema head reads, with lm_head.weight on the host). No lm_head, no table, no state I/O; 1 call = 1 record. |
| |
| Inside (Kev's row graph, kev_work/scripts/kev_graph.py, with embeds and position tables as inputs): |
| mask4d = causal_const + (1 - valid)[:, None, None, :] * -1e4 (causal_const [1, 1, L, L] float32 from a python |
| list: 0 on and below the diagonal, -1e4 above) |
| text_model(inputs_embeds=embeds, attention_mask=mask4d, position_ids=<constant [1, L] buffer>, use_cache=False) |
| A 4-D float mask passes create_causal_mask as is and makes create_recurrent_attention_mask return None; the wrapper |
| sets text_model._clef_valid = valid (stashed on every GatedDeltaNet as `_litert_valid`: the pad guard) and |
| text_model.rotary_emb._clef_cos_sin = (cos, sin) (the patched rotary returns them; the position buffer feeds nothing). |
| |
| Model = transformers 5.14.1's Qwen3_5TextModel built from config.json's text_config with the three classes of |
| scripts/clef_qwen35_patch.py swapped into the modeling module (Kev's order), constructed on the meta device; the |
| `model.language_model.*` tensors of the four shards except embed_tokens (the host's table) are read bf16 -> float32 and |
| assigned in place (`load_state_dict(strict=True, assign=True)`), so the float32 text stack (~27.7 GB) is the only copy. |
| attention = `kev_eager` (GQA by concat, rank-4 BMM, additive mask, float32 softmax) registered in |
| ALL_ATTENTION_FUNCTIONS. The GatedDeltaNet kernel form (`R64+sp+ec+dd+vs6`, scripts/clef_kernel.py) is applied by the |
| caller (`apply_form`). `count_interleave()` proves the ratio-2 head interleave runs (2 calls per GatedDeltaNet).""" |
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
|
|
| K = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(K / "scripts")) |
| import clef_qwen35_patch as P |
| from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| from transformers.models.qwen3_5 import modeling_qwen3_5 as M |
|
|
| REV = "f0b7d3cd65296abfe459e6c228b618f5d8b26d2f" |
| SNAP = K / "hf/hub/models--Cloudflare--clef-flash/snapshots" / REV |
| PREFIX = "model.language_model." |
| NEG = -1e4 |
| ATTN = "kev_eager" |
| FORM = "R64+sp+ec+dd+vs6" |
| HIDDEN = 4096 |
| SWAPS = ("Qwen3_5TextModel", "Qwen3_5GatedDeltaNet", "Qwen3_5TextRotaryEmbedding") |
| ORIG = {name: getattr(M, name) for name in SWAPS} |
| PATCHED = {"Qwen3_5TextModel": P.PatchedQwen3_5TextModel, "Qwen3_5GatedDeltaNet": P.PatchedQwen3_5GatedDeltaNet, |
| "Qwen3_5TextRotaryEmbedding": P.PatchedQwen3_5TextRotaryEmbedding} |
|
|
|
|
| def kev_eager(module, query, key, value, attention_mask, scaling, dropout=0.0, **kwargs): |
| """kev_work/scripts/kev_graph.py's attention, unchanged: key/value [1, Hkv, L, D] -> [Hkv, 1, L, D] -> n_rep copies |
| by concat on dim 1 -> [1, Hkv * n_rep, L, D] (repeat_kv's order), scores = q @ k^T * scaling + mask (rank-4 BMM), |
| softmax in float32, out = p @ v -> [1, L, H, D].""" |
| n_rep = module.num_key_value_groups |
| b, hkv, length, dim = key.shape |
| assert b == 1, "batch 1 only (the batch dim is folded to keep the copy axis inside rank 4)" |
| key_states = torch.cat([key.reshape(hkv, 1, length, dim)] * n_rep, dim=1).reshape(1, hkv * n_rep, length, dim) |
| value_states = torch.cat([value.reshape(hkv, 1, length, dim)] * n_rep, dim=1).reshape(1, hkv * n_rep, length, dim) |
| attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling |
| attn_weights = attn_weights + attention_mask |
| attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) |
| attn_output = torch.matmul(attn_weights, value_states) |
| attn_output = attn_output.transpose(1, 2).contiguous() |
| return attn_output, None |
|
|
|
|
| def count_interleave(): |
| """kev_graph.count_interleave, unchanged: a call counter around the guarded forward's `_litert_interleave_heads` |
| (same function underneath). -> {"calls", "ratios", "shapes"}.""" |
| g = P.PatchedQwen3_5GatedDeltaNet.forward.__globals__ |
| inner = getattr(g["_litert_interleave_heads"], "__wrapped__", g["_litert_interleave_heads"]) |
| counter = {"calls": 0, "ratios": [], "shapes": []} |
|
|
| def counted(x, r): |
| counter["calls"] += 1 |
| if len(counter["ratios"]) < 4: |
| counter["ratios"].append(int(r)) |
| counter["shapes"].append(list(x.shape)) |
| return inner(x, r) |
|
|
| counted.__wrapped__ = inner |
| g["_litert_interleave_heads"] = counted |
| return counter |
|
|
|
|
| def _set_classes(classes: dict) -> None: |
| pkg = sys.modules["transformers.models.qwen3_5"] |
| for name, cls in classes.items(): |
| setattr(M, name, cls) |
| if name in vars(pkg): |
| setattr(pkg, name, cls) |
|
|
|
|
| def causal_constant(L: int) -> torch.Tensor: |
| rows = [[0.0 if j <= i else NEG for j in range(L)] for i in range(L)] |
| return torch.tensor(rows, dtype=torch.float32).reshape(1, 1, L, L) |
|
|
|
|
| def text_config(attn: str = ATTN): |
| from transformers import AutoConfig |
| cfg = AutoConfig.from_pretrained(str(SNAP)).text_config |
| cfg._attn_implementation = attn |
| return cfg |
|
|
|
|
| def _shard_keys(model: nn.Module) -> dict: |
| """{shard file: [module key, ...]} for every tensor of the model's state dict (checkpoint key = PREFIX + key).""" |
| weight_map = json.loads((SNAP / "model.safetensors.index.json").read_text())["weight_map"] |
| ckpt = {k[len(PREFIX):]: f for k, f in weight_map.items() if k.startswith(PREFIX)} |
| assert ckpt.pop("embed_tokens.weight", None) is not None, "no embed_tokens in the checkpoint map" |
| names = set(model.state_dict()) |
| assert names == set(ckpt), (sorted(names - set(ckpt))[:5], sorted(set(ckpt) - names)[:5]) |
| by_shard: dict = {} |
| for k, f in sorted(ckpt.items()): |
| by_shard.setdefault(f, []).append(k) |
| return by_shard |
|
|
|
|
| def _build(cfg, classes: dict) -> nn.Module: |
| _set_classes(classes) |
| try: |
| with torch.device("meta"): |
| model = M.Qwen3_5TextModel(cfg) |
| finally: |
| _set_classes(PATCHED) |
| model.embed_tokens = None |
| return model |
|
|
|
|
| def _finish(model: nn.Module, cfg, rotary_cls) -> nn.Module: |
| model.rotary_emb = rotary_cls(config=cfg) |
| meta = [n for n, t in list(model.named_parameters()) + list(model.named_buffers()) if t.is_meta] |
| assert not meta, meta[:5] |
| return model.eval().requires_grad_(False) |
|
|
|
|
| def load_text_model(): |
| """-> (patched text model, float32, CPU, embed_tokens None; info). Weights read once from the shards.""" |
| from safetensors import safe_open |
| ALL_ATTENTION_FUNCTIONS.register(ATTN, kev_eager) |
| cfg = text_config() |
| t0 = time.perf_counter() |
| model = _build(cfg, PATCHED) |
| by_shard = _shard_keys(model) |
| sd = {} |
| for shard, keys in by_shard.items(): |
| with safe_open(str(SNAP / shard), framework="pt") as f: |
| for k in keys: |
| t = f.get_tensor(PREFIX + k) |
| sd[k] = t.to(torch.float32) |
| src_dtypes = sorted({str(t.dtype) for t in sd.values()}) |
| model.load_state_dict(sd, strict=True, assign=True) |
| del sd |
| model = _finish(model, cfg, P.PatchedQwen3_5TextRotaryEmbedding) |
| load_s = time.perf_counter() - t0 |
| assert type(model) is P.PatchedQwen3_5TextModel, type(model) |
| gdn = [m for m in model.modules() if isinstance(m, M.Qwen3_5GatedDeltaNet)] |
| attn = [m for m in model.modules() if isinstance(m, M.Qwen3_5Attention)] |
| assert gdn and all(type(m) is P.PatchedQwen3_5GatedDeltaNet for m in gdn) |
| assert all(m.chunk_gated_delta_rule is P._rank4_chunk_gated_delta_rule for m in gdn) |
| layer_types = [layer.block_type for layer in model.layers] |
| attn_idx = [i for i, t in enumerate(layer_types) if t == "full_attention"] |
| assert len(model.layers) == 32 and len(gdn) == 24 and len(attn) == 8, (len(model.layers), len(gdn), len(attn)) |
| assert attn_idx == list(range(3, 32, 4)), attn_idx |
| assert (gdn[0].num_k_heads, gdn[0].num_v_heads, gdn[0].head_k_dim, gdn[0].head_v_dim) == (16, 32, 128, 128) |
| assert (cfg.num_attention_heads, cfg.num_key_value_heads, attn[0].head_dim) == (16, 4, 256) |
| assert cfg.hidden_size == HIDDEN and isinstance(model.norm, M.Qwen3_5RMSNorm) |
| assert all(m.config._attn_implementation == ATTN for m in attn) |
| n_params = sum(p.numel() for p in model.parameters()) |
| info = {"model_class": type(model).__name__, "layers": len(model.layers), "gdn_layers": len(gdn), |
| "attention_layers": attn_idx, "attn_implementation": cfg._attn_implementation, |
| "gdn_heads": {"k": gdn[0].num_k_heads, "v": gdn[0].num_v_heads, "ratio": gdn[0].num_v_heads // gdn[0].num_k_heads, |
| "k_dim": gdn[0].head_k_dim, "v_dim": gdn[0].head_v_dim}, |
| "attention_heads": {"q": cfg.num_attention_heads, "kv": cfg.num_key_value_heads, "head_dim": attn[0].head_dim, |
| "num_key_value_groups": attn[0].num_key_value_groups}, |
| "hidden_size": cfg.hidden_size, "intermediate_size": cfg.intermediate_size, "final_norm": type(model.norm).__name__, |
| "chunk_kernel": P._rank4_chunk_gated_delta_rule.__name__, "params": n_params, "fp32_bytes": n_params * 4, |
| "tensors": sum(len(v) for v in by_shard.values()), "source_dtypes": src_dtypes, "embed_tokens": None, |
| "load_seconds": round(load_s, 1), "rope_parameters": cfg.rope_parameters, "rms_norm_eps": cfg.rms_norm_eps} |
| return model, info |
|
|
|
|
| def stock_twin(model: nn.Module) -> nn.Module: |
| """The same weights (shared storage, no copy) in the stock classes of transformers 5.14.1 with sdpa attention: the |
| environment row of the torch check (the stock forward computes its own positions and causal mask).""" |
| cfg = copy.deepcopy(model.config) |
| cfg._attn_implementation = "sdpa" |
| stock = _build(cfg, ORIG) |
| stock.load_state_dict(model.state_dict(), strict=True, assign=True) |
| return _finish(stock, cfg, ORIG["Qwen3_5TextRotaryEmbedding"]) |
|
|
|
|
| def apply_form(model: nn.Module, form: str = FORM) -> dict: |
| import clef_kernel as R |
| return {k: v for k, v in R.apply(model, form).items() if k != "tokens"} |
|
|
|
|
| class ClefRowPrefill(nn.Module): |
| def __init__(self, text_model: nn.Module, L: int): |
| super().__init__() |
| self.text_model = text_model |
| self.L = L |
| self.register_buffer("position_ids", torch.arange(L, dtype=torch.long).reshape(1, L), persistent=False) |
| self.register_buffer("causal_const", causal_constant(L), persistent=False) |
|
|
| def forward(self, embeds, valid, cos, sin): |
| self.text_model._clef_valid = valid |
| self.text_model.rotary_emb._clef_cos_sin = (cos, sin) |
| mask4d = self.causal_const + (1.0 - valid)[:, None, None, :] * NEG |
| out = self.text_model(inputs_embeds=embeds, attention_mask=mask4d, position_ids=self.position_ids, |
| use_cache=False, past_key_values=None) |
| return {"hidden": out.last_hidden_state} |
|
|
|
|
| def graph(model: nn.Module, L: int) -> nn.Module: |
| return ClefRowPrefill(model, L).eval().requires_grad_(False) |
|
|
|
|
| def contract(L: int) -> dict: |
| return { |
| "graph": "ClefRowPrefill (state-free causal row prefill, fp32)", "L": L, "form": FORM, |
| "inputs": [{"name": "embeds", "shape": [1, L, HIDDEN], "dtype": "float32", |
| "meaning": "the checkpoint's bf16 embed_tokens rows as float32; pad rows 0"}, |
| {"name": "valid", "shape": [1, L], "dtype": "float32", "meaning": "1.0 = the record's token, 0.0 = right pad"}, |
| {"name": "cos", "shape": [1, L, 64], "dtype": "float32", "meaning": "M-RoPE cos table (scripts/clef_host.py)"}, |
| {"name": "sin", "shape": [1, L, 64], "dtype": "float32", "meaning": "M-RoPE sin table (scripts/clef_host.py)"}], |
| "outputs": [{"name": "hidden", "shape": [1, L, HIDDEN], "dtype": "float32", |
| "meaning": "last_hidden_state (after the final RMSNorm), every position"}], |
| "pad_rule": "pads only after the record's last token", "attention_mask": "causal_const + (1 - valid) * -1e4", |
| "gated_deltanet_pad_guard": "valid zeroes pad hidden states before the projections, sets pad decay to identity " |
| "(a -> -30 pre-softplus) and re-zeroes pad q/k/v after the causal conv", |
| "state_io": "none (every record starts from zero conv / recurrent state)", |
| } |
|
|
