phanerozoic commited on Aug 11

Commit

418a182

verified ·

1 Parent(s): 2fc8e3b

Uploaded using `kernel-builder`.

Browse files

Files changed (42) hide show

build/torch212-cxx11-cu126-x86_64-linux/__init__.py +23 -0
build/torch212-cxx11-cu126-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so +3 -0
build/torch212-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu126-x86_64-linux/drafter.py +236 -0
build/torch212-cxx11-cu126-x86_64-linux/gguf_io.py +61 -0
build/torch212-cxx11-cu126-x86_64-linux/metadata.json +39 -0
build/torch212-cxx11-cu126-x86_64-linux/reference.py +188 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +23 -0
build/torch212-cxx11-cu130-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/drafter.py +236 -0
build/torch212-cxx11-cu130-x86_64-linux/gguf_io.py +61 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +41 -0
build/torch212-cxx11-cu130-x86_64-linux/reference.py +188 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +23 -0
build/torch212-cxx11-cu132-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/drafter.py +236 -0
build/torch212-cxx11-cu132-x86_64-linux/gguf_io.py +61 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +41 -0
build/torch212-cxx11-cu132-x86_64-linux/reference.py +188 -0
build/torch213-cxx11-cu126-x86_64-linux/__init__.py +23 -0
build/torch213-cxx11-cu126-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so +3 -0
build/torch213-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu126-x86_64-linux/drafter.py +236 -0
build/torch213-cxx11-cu126-x86_64-linux/gguf_io.py +61 -0
build/torch213-cxx11-cu126-x86_64-linux/metadata.json +39 -0
build/torch213-cxx11-cu126-x86_64-linux/reference.py +188 -0
build/torch213-cxx11-cu130-x86_64-linux/__init__.py +23 -0
build/torch213-cxx11-cu130-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so +3 -0
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-x86_64-linux/drafter.py +236 -0
build/torch213-cxx11-cu130-x86_64-linux/gguf_io.py +61 -0
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +41 -0
build/torch213-cxx11-cu130-x86_64-linux/reference.py +188 -0
build/torch213-cxx11-cu132-x86_64-linux/__init__.py +23 -0
build/torch213-cxx11-cu132-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so +3 -0
build/torch213-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu132-x86_64-linux/drafter.py +236 -0
build/torch213-cxx11-cu132-x86_64-linux/gguf_io.py +61 -0
build/torch213-cxx11-cu132-x86_64-linux/metadata.json +41 -0
build/torch213-cxx11-cu132-x86_64-linux/reference.py +188 -0

build/torch212-cxx11-cu126-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""DFlash block-diffusion drafter ops; the eager reference defines the
+semantics and the CUDA ops are certified against it."""
+from .reference import DFlashConfig, DFlashRef, random_weights  # noqa: F401
+from .drafter import DFlashDrafter  # noqa: F401
+try:  # built variants ship _ops
+    from ._ops import ops  # noqa: F401
+except ImportError:
+    try:  # a local JIT build registers the same ops under torch.ops.dflash_ext
+        import torch as _torch
+        _torch.ops.dflash_ext.dflash_accept
+        ops = _torch.ops.dflash_ext
+    except Exception:
+        ops = None
+try:
+    from .gguf_io import load_drafter  # noqa: F401
+except ImportError:  # gguf-py absent; the reference stays usable
+    load_drafter = None
+__all__ = ["DFlashConfig", "DFlashRef", "DFlashDrafter", "random_weights",
+           "load_drafter", "ops"]

build/torch212-cxx11-cu126-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:3451aca743d6ce46a183b58a065863c9abd327abc6a877259861334ea409a6ca
+size 893416

build/torch212-cxx11-cu126-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _dflash_ops_cuda_03c9586
+ops = torch.ops._dflash_ops_cuda_03c9586
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_dflash_ops_cuda_03c9586::{op_name}"

build/torch212-cxx11-cu126-x86_64-linux/drafter.py ADDED Viewed

	@@ -0,0 +1,236 @@

+"""GPU DFlash drafter: cuBLAS GEMMs plus the fused ops, with a CUDA-graphed
+block step."""
+import torch
+from .reference import DFlashConfig, rms_norm, rope_inv_freq
+def _ops():
+    from . import ops
+    if ops is None:
+        raise RuntimeError(
+            "dflash-ops compiled kernels are unavailable; build them with "
+            "load_local.load() in a source checkout, or install a published "
+            "variant with kernels.get_kernel")
+    return ops
+class DFlashDrafter:
+    """Ring-cached DFlash drafter on CUDA. Weight keys match `DFlashRef`."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cuda",
+                 ring_dtype="fp32", dtype="fp32"):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        # dtype sets GEMM weights and activations; residual, norms, softmax stay fp32
+        self.wdt = torch.bfloat16 if dtype == "bf16" else torch.float32
+        self.w = {}
+        for k, v in weights.items():
+            dt = self.wdt if v.dim() == 2 else torch.float32
+            self.w[k] = v.to(device=self.device, dtype=dt).contiguous()
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        # host-built fp64 rope table, shared by every op call
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base,
+                                      self.device).contiguous()
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=self.device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int32,
+                                   device=self.device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        g = (feats.to(self.device, self.wdt) @ self.w["fc"]).float()
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Ring injection at `pos`, chunked at `window` so no slot is written twice per launch."""
+        C = self.cfg
+        o = _ops()
+        pos = pos.to(self.device, torch.int32).contiguous()
+        for beg in range(0, pos.numel(), C.window):
+            p = pos[beg:beg + C.window].contiguous()
+            gc = g[beg:beg + C.window].to(self.wdt)
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                o.dflash_inject(self.ring_k[i], self.ring_v[i], self.ring_pos,
+                                k, v, self.w[f"k_norm.{i}"], self.rope_inv, p,
+                                C.rms_eps)
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = int(id_last)
+        bpos = (n_past + torch.arange(B, device=self.device)).to(torch.int32)
+        return self._forward(tokens, bpos.contiguous())
+    def _forward(self, tokens, bpos):
+        """Block forward on caller-owned device tensors; reads no host state."""
+        C = self.cfg
+        o = _ops()
+        B = tokens.numel()
+        x = self.w["tok_embd"][tokens].float()
+        scale = C.head_dim ** -0.5
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps).to(self.wdt)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head,
+                                             C.head_dim).contiguous()
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            o.dflash_qk_rope(q, k, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, bpos,
+                             C.rms_eps)
+            att = torch.empty_like(q)
+            o.dflash_block_attn(att, q, k, v, self.ring_k[i], self.ring_v[i],
+                                self.ring_pos, bpos, scale)
+            x = x + (att.reshape(B, -1) @ self.w[f"wo.{i}"]).float()
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps).to(self.wdt)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + ((gate * (f @ self.w[f"ffn_up.{i}"]))
+                     @ self.w[f"ffn_down.{i}"]).float()
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return (x.to(self.wdt) @ self.w["lm_head"]).float()
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        return self.block_step(id_last, n_past, n_draft)[1:].argmax(dim=-1)
+    # ---- graphed step ---------------------------------------------------
+    def _static_buffers(self, B):
+        """Every tensor a captured step touches, allocated once and held."""
+        C = self.cfg
+        d = self.device
+        wdt = self.wdt
+        s = {
+            "tokens": torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                                 device=d),
+            "arange": torch.arange(B, dtype=torch.int32, device=d),
+            "n_past": torch.zeros((), dtype=torch.int32, device=d),
+            "bpos": torch.empty(B, dtype=torch.int32, device=d),
+            "xw": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "x": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "sq": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "ms": torch.empty(B, 1, dtype=torch.float32, device=d),
+            "hn": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "h": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "q": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "k": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "v": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "att": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "ao": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "g1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "sg": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "u1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "d1": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "lg": torch.empty(B, C.n_vocab, dtype=wdt, device=d),
+            "logits": torch.empty(B, C.n_vocab, dtype=torch.float32,
+                                  device=d),
+        }
+        return s
+    def _rms_static(self, src, w, s, out):
+        """RMS norm of `src` into `out` using only held buffers."""
+        torch.pow(src, 2, out=s["sq"])
+        torch.mean(s["sq"], -1, keepdim=True, out=s["ms"])
+        s["ms"].add_(self.cfg.rms_eps).rsqrt_()
+        torch.mul(src, s["ms"], out=out)
+        out.mul_(w)
+    def _forward_static(self, s):
+        """Allocation-free block forward into held buffers; numerically
+        identical to `_forward`."""
+        C = self.cfg
+        o = _ops()
+        B = s["tokens"].numel()
+        scale = C.head_dim ** -0.5
+        torch.index_select(self.w["tok_embd"], 0, s["tokens"], out=s["xw"])
+        s["x"].copy_(s["xw"])
+        for i in range(C.n_layer):
+            self._rms_static(s["x"], self.w[f"attn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"wq.{i}"], out=s["q"])
+            torch.matmul(s["h"], self.w[f"wk.{i}"], out=s["k"])
+            torch.matmul(s["h"], self.w[f"wv.{i}"], out=s["v"])
+            qv = s["q"].view(B, C.n_head, C.head_dim)
+            kv_ = s["k"].view(B, C.n_head_kv, C.head_dim)
+            vv = s["v"].view(B, C.n_head_kv, C.head_dim)
+            o.dflash_qk_rope(qv, kv_, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, s["bpos"],
+                             C.rms_eps)
+            av = s["att"].view(B, C.n_head, C.head_dim)
+            o.dflash_block_attn(av, qv, kv_, vv, self.ring_k[i],
+                                self.ring_v[i], self.ring_pos, s["bpos"],
+                                scale)
+            torch.matmul(s["att"], self.w[f"wo.{i}"], out=s["ao"])
+            s["x"].add_(s["ao"])
+            self._rms_static(s["x"], self.w[f"ffn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"ffn_gate.{i}"], out=s["g1"])
+            # fused silu keeps the step bitwise-equal to the eager path
+            torch.ops.aten.silu.out(s["g1"], out=s["sg"])
+            torch.matmul(s["h"], self.w[f"ffn_up.{i}"], out=s["u1"])
+            s["sg"].mul_(s["u1"])
+            torch.matmul(s["sg"], self.w[f"ffn_down.{i}"], out=s["d1"])
+            s["x"].add_(s["d1"])
+        self._rms_static(s["x"], self.w["out_norm"], s, s["hn"])
+        s["h"].copy_(s["hn"])
+        torch.matmul(s["h"], self.w["lm_head"], out=s["lg"])
+        s["logits"].copy_(s["lg"])
+        return s["logits"]
+    def capture_step(self, n_draft=None):
+        """Graph-captured block step over static buffers (the region allocates
+        nothing); returns step(id_last, n_past) -> logits, overwritten per replay."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        s = self._static_buffers(n_draft + 1)
+        stream = torch.cuda.Stream(device=self.device)
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):                    # warm kernels and cuBLAS
+                torch.add(s["arange"], s["n_past"], out=s["bpos"])
+                self._forward_static(s)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            torch.add(s["arange"], s["n_past"], out=s["bpos"])  # in-graph: replay sees updates
+            logits = self._forward_static(s)
+        def step(id_last, n_past):
+            s["tokens"][0].fill_(int(id_last))
+            s["n_past"].fill_(int(n_past))
+            graph.replay()
+            return logits
+        step.graph = graph
+        step.buffers = s
+        return step
+    # ---- verification ---------------------------------------------------
+    def accept(self, draft, target):
+        """Leading agreement count, device-side."""
+        o = _ops()
+        n = torch.zeros(1, dtype=torch.int32, device=self.device)
+        o.dflash_accept(n, draft.to(torch.int32).contiguous(),
+                        target.to(torch.int32).contiguous())
+        return n

build/torch212-cxx11-cu126-x86_64-linux/gguf_io.py ADDED Viewed

	@@ -0,0 +1,61 @@

+"""Load the DFlash drafter GGUF into reference-layout fp32 tensors; 2-D weights
+transpose from GGUF [out, in] to [in, out]; tok_embd and lm_head are the
+caller's (they belong to the target model)."""
+import numpy as np
+import torch
+from gguf import GGUFReader
+from gguf import quants as gq
+from .reference import DFlashConfig
+_BLK = {
+    "attn_norm": "attn_norm", "attn_q": "wq", "attn_k": "wk", "attn_v": "wv",
+    "attn_output": "wo", "attn_q_norm": "q_norm", "attn_k_norm": "k_norm",
+    "ffn_norm": "ffn_norm", "ffn_gate": "ffn_gate", "ffn_up": "ffn_up",
+    "ffn_down": "ffn_down",
+}
+_TOP = {"fc": "fc", "enc.output_norm": "enc_norm", "output_norm": "out_norm"}
+def _field(reader, key, default=None):
+    f = reader.get_field(key)
+    if f is None:
+        return default
+    v = f.contents()
+    return v
+def load_drafter(path):
+    """-> (DFlashConfig, weights dict) with fp32 tensors in reference layout."""
+    r = GGUFReader(str(path))
+    cfg = DFlashConfig(
+        n_layer=int(_field(r, "dflash.block_count")),
+        n_embd=int(_field(r, "dflash.embedding_length")),
+        n_head=int(_field(r, "dflash.attention.head_count")),
+        n_head_kv=int(_field(r, "dflash.attention.head_count_kv")),
+        head_dim=int(_field(r, "dflash.attention.key_length")),
+        n_ff=int(_field(r, "dflash.feed_forward_length")),
+        rope_base=float(_field(r, "dflash.rope.freq_base")),
+        rms_eps=float(_field(r, "dflash.attention.layer_norm_rms_epsilon")),
+        window=int(_field(r, "dflash.attention.sliding_window")),
+        block_size=int(_field(r, "dflash.block_size")),
+        target_layers=tuple(_field(r, "dflash.target_layers")),
+    )
+    w = {}
+    for t in r.tensors:
+        name = t.name
+        assert name.endswith(".weight"), name
+        base = name[: -len(".weight")]
+        if base.startswith("blk."):
+            _, idx, part = base.split(".", 2)
+            key = f"{_BLK[part]}.{idx}"
+        else:
+            key = _TOP[base]
+        data = gq.dequantize(t.data, t.tensor_type)
+        arr = torch.from_numpy(np.ascontiguousarray(data)).float()
+        if arr.dim() == 2:
+            arr = arr.t().contiguous()
+        w[key] = arr
+    expect = 3 + 11 * cfg.n_layer
+    assert len(w) == expect, (len(w), expect)
+    return cfg, w

build/torch212-cxx11-cu126-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "dflash-ops",
+  "id": "_dflash_ops_cuda_03c9586",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "wNpuIbQEXTfFNjrmz6aISM46VFgMKAHkqqu+ECJJS20=",
+      "_dflash_ops_cuda_03c9586.abi3.so": "NFGsp0PWzkahg7WKBlhjyavTJ6vGqHclmGEzTqQJpso=",
+      "_ops.py": "+6e3Oo6lwRkZaeOrhSSZEwWYs4/Fpd/zQgJj4MVEi9I=",
+      "drafter.py": "wGU5TWlXWr/ZogxQr4pho5Wku66UUygnczpsTu82dls=",
+      "gguf_io.py": "XrgF2dH7LNR1KTNKKu/pKhh6wT90bCjcjCAgmyGsuZA=",
+      "reference.py": "InmFS9rUwURZ7ZGVUxwQJCCjJZR/hKWrF3eMfp4zQvY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "03c9586154cc2afa6efaa7b111dccef715ec9d69",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu126-x86_64-linux/reference.py ADDED Viewed

	@@ -0,0 +1,188 @@

+"""Eager fp32 DFlash drafter reference, transcribed from llama.cpp's
+llama_model_dflash graph; the CUDA ops are certified against this module."""
+from dataclasses import dataclass, field
+import torch
+@dataclass
+class DFlashConfig:
+    n_layer: int = 5
+    n_embd: int = 6656
+    n_head: int = 32
+    n_head_kv: int = 8
+    head_dim: int = 128
+    n_ff: int = 19968
+    n_vocab: int = 202048
+    rope_base: float = 500000.0
+    rms_eps: float = 1e-6
+    window: int = 2048            # SWA width; also the ring capacity
+    block_size: int = 16
+    mask_token_id: int = 201818
+    target_layers: tuple = (2, 14, 26, 38, 50)
+    @property
+    def n_feat(self):
+        return len(self.target_layers) * self.n_embd
+    @classmethod
+    def tiny(cls):
+        """Toy sizes with the real structure; head_dim stays a multiple of 32."""
+        return cls(n_layer=2, n_embd=64, n_head=4, n_head_kv=2, head_dim=32,
+                   n_ff=96, n_vocab=997, window=32, block_size=8,
+                   mask_token_id=990, target_layers=(1, 2))
+def rms_norm(x, w, eps):
+    """llama.cpp LLM_NORM_RMS: x * w / sqrt(mean(x^2) + eps), fp32."""
+    v = x.float()
+    return (v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps)) * w.float()
+def rope_inv_freq(head_dim, base, device="cpu"):
+    """fp64 NEOX inverse frequencies base^(-2i/head_dim), host-evaluated so the
+    table is device-invariant."""
+    half = head_dim // 2
+    t = base ** (-torch.arange(half, dtype=torch.float64) * 2.0 / head_dim)
+    return t.to(device)
+def rope_neox(x, pos, base=None, inv=None):
+    """NEOX rope over pairs (i, i + d/2), fp32; x [T, H, D], pos [T]; pass
+    `inv` to reuse a precomputed table."""
+    T, H, D = x.shape
+    half = D // 2
+    if inv is None:
+        inv = rope_inv_freq(D, base, x.device)
+    ang = pos.to(torch.float64)[:, None] * inv[None, :]          # [T, half]
+    cos = ang.cos().to(torch.float32)[:, None, :]
+    sin = ang.sin().to(torch.float32)[:, None, :]
+    xf = x.float()
+    x1, x2 = xf[..., :half], xf[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+class DFlashRef:
+    """Reference drafter over an fp32 weight dict (layout as in
+    `random_weights`) with a sliding ring KV state."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cpu",
+                 ring_dtype="fp32"):
+        self.cfg = cfg
+        self.w = {k: v.to(device=device, dtype=torch.float32)
+                  for k, v in weights.items()}
+        self.device = torch.device(device)
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base, device)
+        # per-layer ring K/V [window, Hkv, D]; ring_pos = slot's position, -1 empty
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int64,
+                                   device=device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        """Fuse per-position target features [T, n_feat] to g [T, n_embd]."""
+        g = feats.float() @ self.w["fc"]
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Project fused features into every layer's ring at `pos` [T], chunked
+        at `window` so no two tokens in one pass share a slot."""
+        C = self.cfg
+        for beg in range(0, len(pos), C.window):
+            p = pos[beg:beg + C.window]
+            gc = g[beg:beg + C.window]
+            slots = p % C.window
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+                k = rope_neox(k, p, inv=self.rope_inv)
+                self.ring_k[i, slots] = k.to(self.ring_dt)
+                self.ring_v[i, slots] = v.to(self.ring_dt)
+            self.ring_pos[slots] = p
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        """Denoise [id_last, mask x n_draft] at n_past..; -> [n_draft+1, n_vocab]."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = id_last
+        pos = n_past + torch.arange(B, device=self.device)
+        x = self.w["tok_embd"][tokens]
+        live = self.ring_pos >= 0
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head, C.head_dim)
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            q = rms_norm(q, self.w[f"q_norm.{i}"], C.rms_eps)
+            k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+            q = rope_neox(q, pos, inv=self.rope_inv)
+            k = rope_neox(k, pos, inv=self.rope_inv)
+            # keys = live ring + block; non-causal, SWA masks only the deep past
+            keys = torch.cat([self.ring_k[i, live].float(), k], dim=0)
+            vals = torch.cat([self.ring_v[i, live].float(), v], dim=0)
+            kpos = torch.cat([self.ring_pos[live], pos], dim=0)
+            gq = C.n_head // C.n_head_kv
+            qh = q.permute(1, 0, 2)                                # [H, B, D]
+            kh = keys.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            vh = vals.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            att = qh @ kh.transpose(1, 2) / (C.head_dim ** 0.5)    # [H, B, S]
+            masked = (pos[:, None] - kpos[None, :]) >= C.window    # SWA
+            att = att.masked_fill(masked[None, :, :], float("-inf"))
+            att = att.softmax(dim=-1)
+            out = (att @ vh).permute(1, 0, 2).reshape(B, -1)       # [B, H*D]
+            x = x + out @ self.w[f"wo.{i}"]
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + (gate * (f @ self.w[f"ffn_up.{i}"])) @ self.w[f"ffn_down.{i}"]
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return x @ self.w["lm_head"]
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        """Greedy drafts from block positions 1..n_draft (position 0 unused)."""
+        logits = self.block_step(id_last, n_past, n_draft)
+        return logits[1:].argmax(dim=-1)
+def random_weights(cfg: DFlashConfig, seed=0, device="cpu"):
+    """Gaussian weights at transformer-typical scale, for structural tests."""
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    def t(*shape, s=0.02):
+        return torch.randn(*shape, generator=g) * s
+    C = cfg
+    w = {"fc": t(C.n_feat, C.n_embd), "enc_norm": torch.ones(C.n_embd),
+         "out_norm": torch.ones(C.n_embd),
+         "tok_embd": t(C.n_vocab, C.n_embd),
+         "lm_head": t(C.n_embd, C.n_vocab)}
+    for i in range(C.n_layer):
+        w[f"attn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"wq.{i}"] = t(C.n_embd, C.n_head * C.head_dim)
+        w[f"wk.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wv.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wo.{i}"] = t(C.n_head * C.head_dim, C.n_embd)
+        w[f"q_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"k_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"ffn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"ffn_gate.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_up.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_down.{i}"] = t(C.n_ff, C.n_embd)
+    return {k: v.to(device) for k, v in w.items()}

build/torch212-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""DFlash block-diffusion drafter ops; the eager reference defines the
+semantics and the CUDA ops are certified against it."""
+from .reference import DFlashConfig, DFlashRef, random_weights  # noqa: F401
+from .drafter import DFlashDrafter  # noqa: F401
+try:  # built variants ship _ops
+    from ._ops import ops  # noqa: F401
+except ImportError:
+    try:  # a local JIT build registers the same ops under torch.ops.dflash_ext
+        import torch as _torch
+        _torch.ops.dflash_ext.dflash_accept
+        ops = _torch.ops.dflash_ext
+    except Exception:
+        ops = None
+try:
+    from .gguf_io import load_drafter  # noqa: F401
+except ImportError:  # gguf-py absent; the reference stays usable
+    load_drafter = None
+__all__ = ["DFlashConfig", "DFlashRef", "DFlashDrafter", "random_weights",
+           "load_drafter", "ops"]

build/torch212-cxx11-cu130-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:2f92dc3409c1896355af5a96649ff74e918c3d460ea0ea7c73080bd121e75085
+size 1402256

build/torch212-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _dflash_ops_cuda_03c9586
+ops = torch.ops._dflash_ops_cuda_03c9586
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_dflash_ops_cuda_03c9586::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/drafter.py ADDED Viewed

	@@ -0,0 +1,236 @@

+"""GPU DFlash drafter: cuBLAS GEMMs plus the fused ops, with a CUDA-graphed
+block step."""
+import torch
+from .reference import DFlashConfig, rms_norm, rope_inv_freq
+def _ops():
+    from . import ops
+    if ops is None:
+        raise RuntimeError(
+            "dflash-ops compiled kernels are unavailable; build them with "
+            "load_local.load() in a source checkout, or install a published "
+            "variant with kernels.get_kernel")
+    return ops
+class DFlashDrafter:
+    """Ring-cached DFlash drafter on CUDA. Weight keys match `DFlashRef`."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cuda",
+                 ring_dtype="fp32", dtype="fp32"):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        # dtype sets GEMM weights and activations; residual, norms, softmax stay fp32
+        self.wdt = torch.bfloat16 if dtype == "bf16" else torch.float32
+        self.w = {}
+        for k, v in weights.items():
+            dt = self.wdt if v.dim() == 2 else torch.float32
+            self.w[k] = v.to(device=self.device, dtype=dt).contiguous()
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        # host-built fp64 rope table, shared by every op call
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base,
+                                      self.device).contiguous()
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=self.device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int32,
+                                   device=self.device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        g = (feats.to(self.device, self.wdt) @ self.w["fc"]).float()
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Ring injection at `pos`, chunked at `window` so no slot is written twice per launch."""
+        C = self.cfg
+        o = _ops()
+        pos = pos.to(self.device, torch.int32).contiguous()
+        for beg in range(0, pos.numel(), C.window):
+            p = pos[beg:beg + C.window].contiguous()
+            gc = g[beg:beg + C.window].to(self.wdt)
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                o.dflash_inject(self.ring_k[i], self.ring_v[i], self.ring_pos,
+                                k, v, self.w[f"k_norm.{i}"], self.rope_inv, p,
+                                C.rms_eps)
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = int(id_last)
+        bpos = (n_past + torch.arange(B, device=self.device)).to(torch.int32)
+        return self._forward(tokens, bpos.contiguous())
+    def _forward(self, tokens, bpos):
+        """Block forward on caller-owned device tensors; reads no host state."""
+        C = self.cfg
+        o = _ops()
+        B = tokens.numel()
+        x = self.w["tok_embd"][tokens].float()
+        scale = C.head_dim ** -0.5
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps).to(self.wdt)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head,
+                                             C.head_dim).contiguous()
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            o.dflash_qk_rope(q, k, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, bpos,
+                             C.rms_eps)
+            att = torch.empty_like(q)
+            o.dflash_block_attn(att, q, k, v, self.ring_k[i], self.ring_v[i],
+                                self.ring_pos, bpos, scale)
+            x = x + (att.reshape(B, -1) @ self.w[f"wo.{i}"]).float()
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps).to(self.wdt)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + ((gate * (f @ self.w[f"ffn_up.{i}"]))
+                     @ self.w[f"ffn_down.{i}"]).float()
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return (x.to(self.wdt) @ self.w["lm_head"]).float()
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        return self.block_step(id_last, n_past, n_draft)[1:].argmax(dim=-1)
+    # ---- graphed step ---------------------------------------------------
+    def _static_buffers(self, B):
+        """Every tensor a captured step touches, allocated once and held."""
+        C = self.cfg
+        d = self.device
+        wdt = self.wdt
+        s = {
+            "tokens": torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                                 device=d),
+            "arange": torch.arange(B, dtype=torch.int32, device=d),
+            "n_past": torch.zeros((), dtype=torch.int32, device=d),
+            "bpos": torch.empty(B, dtype=torch.int32, device=d),
+            "xw": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "x": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "sq": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "ms": torch.empty(B, 1, dtype=torch.float32, device=d),
+            "hn": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "h": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "q": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "k": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "v": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "att": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "ao": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "g1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "sg": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "u1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "d1": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "lg": torch.empty(B, C.n_vocab, dtype=wdt, device=d),
+            "logits": torch.empty(B, C.n_vocab, dtype=torch.float32,
+                                  device=d),
+        }
+        return s
+    def _rms_static(self, src, w, s, out):
+        """RMS norm of `src` into `out` using only held buffers."""
+        torch.pow(src, 2, out=s["sq"])
+        torch.mean(s["sq"], -1, keepdim=True, out=s["ms"])
+        s["ms"].add_(self.cfg.rms_eps).rsqrt_()
+        torch.mul(src, s["ms"], out=out)
+        out.mul_(w)
+    def _forward_static(self, s):
+        """Allocation-free block forward into held buffers; numerically
+        identical to `_forward`."""
+        C = self.cfg
+        o = _ops()
+        B = s["tokens"].numel()
+        scale = C.head_dim ** -0.5
+        torch.index_select(self.w["tok_embd"], 0, s["tokens"], out=s["xw"])
+        s["x"].copy_(s["xw"])
+        for i in range(C.n_layer):
+            self._rms_static(s["x"], self.w[f"attn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"wq.{i}"], out=s["q"])
+            torch.matmul(s["h"], self.w[f"wk.{i}"], out=s["k"])
+            torch.matmul(s["h"], self.w[f"wv.{i}"], out=s["v"])
+            qv = s["q"].view(B, C.n_head, C.head_dim)
+            kv_ = s["k"].view(B, C.n_head_kv, C.head_dim)
+            vv = s["v"].view(B, C.n_head_kv, C.head_dim)
+            o.dflash_qk_rope(qv, kv_, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, s["bpos"],
+                             C.rms_eps)
+            av = s["att"].view(B, C.n_head, C.head_dim)
+            o.dflash_block_attn(av, qv, kv_, vv, self.ring_k[i],
+                                self.ring_v[i], self.ring_pos, s["bpos"],
+                                scale)
+            torch.matmul(s["att"], self.w[f"wo.{i}"], out=s["ao"])
+            s["x"].add_(s["ao"])
+            self._rms_static(s["x"], self.w[f"ffn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"ffn_gate.{i}"], out=s["g1"])
+            # fused silu keeps the step bitwise-equal to the eager path
+            torch.ops.aten.silu.out(s["g1"], out=s["sg"])
+            torch.matmul(s["h"], self.w[f"ffn_up.{i}"], out=s["u1"])
+            s["sg"].mul_(s["u1"])
+            torch.matmul(s["sg"], self.w[f"ffn_down.{i}"], out=s["d1"])
+            s["x"].add_(s["d1"])
+        self._rms_static(s["x"], self.w["out_norm"], s, s["hn"])
+        s["h"].copy_(s["hn"])
+        torch.matmul(s["h"], self.w["lm_head"], out=s["lg"])
+        s["logits"].copy_(s["lg"])
+        return s["logits"]
+    def capture_step(self, n_draft=None):
+        """Graph-captured block step over static buffers (the region allocates
+        nothing); returns step(id_last, n_past) -> logits, overwritten per replay."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        s = self._static_buffers(n_draft + 1)
+        stream = torch.cuda.Stream(device=self.device)
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):                    # warm kernels and cuBLAS
+                torch.add(s["arange"], s["n_past"], out=s["bpos"])
+                self._forward_static(s)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            torch.add(s["arange"], s["n_past"], out=s["bpos"])  # in-graph: replay sees updates
+            logits = self._forward_static(s)
+        def step(id_last, n_past):
+            s["tokens"][0].fill_(int(id_last))
+            s["n_past"].fill_(int(n_past))
+            graph.replay()
+            return logits
+        step.graph = graph
+        step.buffers = s
+        return step
+    # ---- verification ---------------------------------------------------
+    def accept(self, draft, target):
+        """Leading agreement count, device-side."""
+        o = _ops()
+        n = torch.zeros(1, dtype=torch.int32, device=self.device)
+        o.dflash_accept(n, draft.to(torch.int32).contiguous(),
+                        target.to(torch.int32).contiguous())
+        return n

build/torch212-cxx11-cu130-x86_64-linux/gguf_io.py ADDED Viewed

	@@ -0,0 +1,61 @@

+"""Load the DFlash drafter GGUF into reference-layout fp32 tensors; 2-D weights
+transpose from GGUF [out, in] to [in, out]; tok_embd and lm_head are the
+caller's (they belong to the target model)."""
+import numpy as np
+import torch
+from gguf import GGUFReader
+from gguf import quants as gq
+from .reference import DFlashConfig
+_BLK = {
+    "attn_norm": "attn_norm", "attn_q": "wq", "attn_k": "wk", "attn_v": "wv",
+    "attn_output": "wo", "attn_q_norm": "q_norm", "attn_k_norm": "k_norm",
+    "ffn_norm": "ffn_norm", "ffn_gate": "ffn_gate", "ffn_up": "ffn_up",
+    "ffn_down": "ffn_down",
+}
+_TOP = {"fc": "fc", "enc.output_norm": "enc_norm", "output_norm": "out_norm"}
+def _field(reader, key, default=None):
+    f = reader.get_field(key)
+    if f is None:
+        return default
+    v = f.contents()
+    return v
+def load_drafter(path):
+    """-> (DFlashConfig, weights dict) with fp32 tensors in reference layout."""
+    r = GGUFReader(str(path))
+    cfg = DFlashConfig(
+        n_layer=int(_field(r, "dflash.block_count")),
+        n_embd=int(_field(r, "dflash.embedding_length")),
+        n_head=int(_field(r, "dflash.attention.head_count")),
+        n_head_kv=int(_field(r, "dflash.attention.head_count_kv")),
+        head_dim=int(_field(r, "dflash.attention.key_length")),
+        n_ff=int(_field(r, "dflash.feed_forward_length")),
+        rope_base=float(_field(r, "dflash.rope.freq_base")),
+        rms_eps=float(_field(r, "dflash.attention.layer_norm_rms_epsilon")),
+        window=int(_field(r, "dflash.attention.sliding_window")),
+        block_size=int(_field(r, "dflash.block_size")),
+        target_layers=tuple(_field(r, "dflash.target_layers")),
+    )
+    w = {}
+    for t in r.tensors:
+        name = t.name
+        assert name.endswith(".weight"), name
+        base = name[: -len(".weight")]
+        if base.startswith("blk."):
+            _, idx, part = base.split(".", 2)
+            key = f"{_BLK[part]}.{idx}"
+        else:
+            key = _TOP[base]
+        data = gq.dequantize(t.data, t.tensor_type)
+        arr = torch.from_numpy(np.ascontiguousarray(data)).float()
+        if arr.dim() == 2:
+            arr = arr.t().contiguous()
+        w[key] = arr
+    expect = 3 + 11 * cfg.n_layer
+    assert len(w) == expect, (len(w), expect)
+    return cfg, w

build/torch212-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "dflash-ops",
+  "id": "_dflash_ops_cuda_03c9586",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "wNpuIbQEXTfFNjrmz6aISM46VFgMKAHkqqu+ECJJS20=",
+      "_dflash_ops_cuda_03c9586.abi3.so": "L5LcNAnBiWNVr1qWZJ/3TpGMPUYOoOp8cwgL0SHnUIU=",
+      "_ops.py": "+6e3Oo6lwRkZaeOrhSSZEwWYs4/Fpd/zQgJj4MVEi9I=",
+      "drafter.py": "wGU5TWlXWr/ZogxQr4pho5Wku66UUygnczpsTu82dls=",
+      "gguf_io.py": "XrgF2dH7LNR1KTNKKu/pKhh6wT90bCjcjCAgmyGsuZA=",
+      "reference.py": "InmFS9rUwURZ7ZGVUxwQJCCjJZR/hKWrF3eMfp4zQvY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "03c9586154cc2afa6efaa7b111dccef715ec9d69",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/reference.py ADDED Viewed

	@@ -0,0 +1,188 @@

+"""Eager fp32 DFlash drafter reference, transcribed from llama.cpp's
+llama_model_dflash graph; the CUDA ops are certified against this module."""
+from dataclasses import dataclass, field
+import torch
+@dataclass
+class DFlashConfig:
+    n_layer: int = 5
+    n_embd: int = 6656
+    n_head: int = 32
+    n_head_kv: int = 8
+    head_dim: int = 128
+    n_ff: int = 19968
+    n_vocab: int = 202048
+    rope_base: float = 500000.0
+    rms_eps: float = 1e-6
+    window: int = 2048            # SWA width; also the ring capacity
+    block_size: int = 16
+    mask_token_id: int = 201818
+    target_layers: tuple = (2, 14, 26, 38, 50)
+    @property
+    def n_feat(self):
+        return len(self.target_layers) * self.n_embd
+    @classmethod
+    def tiny(cls):
+        """Toy sizes with the real structure; head_dim stays a multiple of 32."""
+        return cls(n_layer=2, n_embd=64, n_head=4, n_head_kv=2, head_dim=32,
+                   n_ff=96, n_vocab=997, window=32, block_size=8,
+                   mask_token_id=990, target_layers=(1, 2))
+def rms_norm(x, w, eps):
+    """llama.cpp LLM_NORM_RMS: x * w / sqrt(mean(x^2) + eps), fp32."""
+    v = x.float()
+    return (v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps)) * w.float()
+def rope_inv_freq(head_dim, base, device="cpu"):
+    """fp64 NEOX inverse frequencies base^(-2i/head_dim), host-evaluated so the
+    table is device-invariant."""
+    half = head_dim // 2
+    t = base ** (-torch.arange(half, dtype=torch.float64) * 2.0 / head_dim)
+    return t.to(device)
+def rope_neox(x, pos, base=None, inv=None):
+    """NEOX rope over pairs (i, i + d/2), fp32; x [T, H, D], pos [T]; pass
+    `inv` to reuse a precomputed table."""
+    T, H, D = x.shape
+    half = D // 2
+    if inv is None:
+        inv = rope_inv_freq(D, base, x.device)
+    ang = pos.to(torch.float64)[:, None] * inv[None, :]          # [T, half]
+    cos = ang.cos().to(torch.float32)[:, None, :]
+    sin = ang.sin().to(torch.float32)[:, None, :]
+    xf = x.float()
+    x1, x2 = xf[..., :half], xf[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+class DFlashRef:
+    """Reference drafter over an fp32 weight dict (layout as in
+    `random_weights`) with a sliding ring KV state."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cpu",
+                 ring_dtype="fp32"):
+        self.cfg = cfg
+        self.w = {k: v.to(device=device, dtype=torch.float32)
+                  for k, v in weights.items()}
+        self.device = torch.device(device)
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base, device)
+        # per-layer ring K/V [window, Hkv, D]; ring_pos = slot's position, -1 empty
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int64,
+                                   device=device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        """Fuse per-position target features [T, n_feat] to g [T, n_embd]."""
+        g = feats.float() @ self.w["fc"]
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Project fused features into every layer's ring at `pos` [T], chunked
+        at `window` so no two tokens in one pass share a slot."""
+        C = self.cfg
+        for beg in range(0, len(pos), C.window):
+            p = pos[beg:beg + C.window]
+            gc = g[beg:beg + C.window]
+            slots = p % C.window
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+                k = rope_neox(k, p, inv=self.rope_inv)
+                self.ring_k[i, slots] = k.to(self.ring_dt)
+                self.ring_v[i, slots] = v.to(self.ring_dt)
+            self.ring_pos[slots] = p
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        """Denoise [id_last, mask x n_draft] at n_past..; -> [n_draft+1, n_vocab]."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = id_last
+        pos = n_past + torch.arange(B, device=self.device)
+        x = self.w["tok_embd"][tokens]
+        live = self.ring_pos >= 0
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head, C.head_dim)
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            q = rms_norm(q, self.w[f"q_norm.{i}"], C.rms_eps)
+            k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+            q = rope_neox(q, pos, inv=self.rope_inv)
+            k = rope_neox(k, pos, inv=self.rope_inv)
+            # keys = live ring + block; non-causal, SWA masks only the deep past
+            keys = torch.cat([self.ring_k[i, live].float(), k], dim=0)
+            vals = torch.cat([self.ring_v[i, live].float(), v], dim=0)
+            kpos = torch.cat([self.ring_pos[live], pos], dim=0)
+            gq = C.n_head // C.n_head_kv
+            qh = q.permute(1, 0, 2)                                # [H, B, D]
+            kh = keys.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            vh = vals.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            att = qh @ kh.transpose(1, 2) / (C.head_dim ** 0.5)    # [H, B, S]
+            masked = (pos[:, None] - kpos[None, :]) >= C.window    # SWA
+            att = att.masked_fill(masked[None, :, :], float("-inf"))
+            att = att.softmax(dim=-1)
+            out = (att @ vh).permute(1, 0, 2).reshape(B, -1)       # [B, H*D]
+            x = x + out @ self.w[f"wo.{i}"]
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + (gate * (f @ self.w[f"ffn_up.{i}"])) @ self.w[f"ffn_down.{i}"]
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return x @ self.w["lm_head"]
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        """Greedy drafts from block positions 1..n_draft (position 0 unused)."""
+        logits = self.block_step(id_last, n_past, n_draft)
+        return logits[1:].argmax(dim=-1)
+def random_weights(cfg: DFlashConfig, seed=0, device="cpu"):
+    """Gaussian weights at transformer-typical scale, for structural tests."""
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    def t(*shape, s=0.02):
+        return torch.randn(*shape, generator=g) * s
+    C = cfg
+    w = {"fc": t(C.n_feat, C.n_embd), "enc_norm": torch.ones(C.n_embd),
+         "out_norm": torch.ones(C.n_embd),
+         "tok_embd": t(C.n_vocab, C.n_embd),
+         "lm_head": t(C.n_embd, C.n_vocab)}
+    for i in range(C.n_layer):
+        w[f"attn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"wq.{i}"] = t(C.n_embd, C.n_head * C.head_dim)
+        w[f"wk.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wv.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wo.{i}"] = t(C.n_head * C.head_dim, C.n_embd)
+        w[f"q_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"k_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"ffn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"ffn_gate.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_up.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_down.{i}"] = t(C.n_ff, C.n_embd)
+    return {k: v.to(device) for k, v in w.items()}

build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""DFlash block-diffusion drafter ops; the eager reference defines the
+semantics and the CUDA ops are certified against it."""
+from .reference import DFlashConfig, DFlashRef, random_weights  # noqa: F401
+from .drafter import DFlashDrafter  # noqa: F401
+try:  # built variants ship _ops
+    from ._ops import ops  # noqa: F401
+except ImportError:
+    try:  # a local JIT build registers the same ops under torch.ops.dflash_ext
+        import torch as _torch
+        _torch.ops.dflash_ext.dflash_accept
+        ops = _torch.ops.dflash_ext
+    except Exception:
+        ops = None
+try:
+    from .gguf_io import load_drafter  # noqa: F401
+except ImportError:  # gguf-py absent; the reference stays usable
+    load_drafter = None
+__all__ = ["DFlashConfig", "DFlashRef", "DFlashDrafter", "random_weights",
+           "load_drafter", "ops"]

build/torch212-cxx11-cu132-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:cd91be1764fb31dd2579c7f2e9f079227873de5405a162b14324ad4b35cc1f35
+size 1402256

build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _dflash_ops_cuda_03c9586
+ops = torch.ops._dflash_ops_cuda_03c9586
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_dflash_ops_cuda_03c9586::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/drafter.py ADDED Viewed

	@@ -0,0 +1,236 @@

+"""GPU DFlash drafter: cuBLAS GEMMs plus the fused ops, with a CUDA-graphed
+block step."""
+import torch
+from .reference import DFlashConfig, rms_norm, rope_inv_freq
+def _ops():
+    from . import ops
+    if ops is None:
+        raise RuntimeError(
+            "dflash-ops compiled kernels are unavailable; build them with "
+            "load_local.load() in a source checkout, or install a published "
+            "variant with kernels.get_kernel")
+    return ops
+class DFlashDrafter:
+    """Ring-cached DFlash drafter on CUDA. Weight keys match `DFlashRef`."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cuda",
+                 ring_dtype="fp32", dtype="fp32"):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        # dtype sets GEMM weights and activations; residual, norms, softmax stay fp32
+        self.wdt = torch.bfloat16 if dtype == "bf16" else torch.float32
+        self.w = {}
+        for k, v in weights.items():
+            dt = self.wdt if v.dim() == 2 else torch.float32
+            self.w[k] = v.to(device=self.device, dtype=dt).contiguous()
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        # host-built fp64 rope table, shared by every op call
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base,
+                                      self.device).contiguous()
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=self.device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int32,
+                                   device=self.device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        g = (feats.to(self.device, self.wdt) @ self.w["fc"]).float()
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Ring injection at `pos`, chunked at `window` so no slot is written twice per launch."""
+        C = self.cfg
+        o = _ops()
+        pos = pos.to(self.device, torch.int32).contiguous()
+        for beg in range(0, pos.numel(), C.window):
+            p = pos[beg:beg + C.window].contiguous()
+            gc = g[beg:beg + C.window].to(self.wdt)
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                o.dflash_inject(self.ring_k[i], self.ring_v[i], self.ring_pos,
+                                k, v, self.w[f"k_norm.{i}"], self.rope_inv, p,
+                                C.rms_eps)
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = int(id_last)
+        bpos = (n_past + torch.arange(B, device=self.device)).to(torch.int32)
+        return self._forward(tokens, bpos.contiguous())
+    def _forward(self, tokens, bpos):
+        """Block forward on caller-owned device tensors; reads no host state."""
+        C = self.cfg
+        o = _ops()
+        B = tokens.numel()
+        x = self.w["tok_embd"][tokens].float()
+        scale = C.head_dim ** -0.5
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps).to(self.wdt)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head,
+                                             C.head_dim).contiguous()
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            o.dflash_qk_rope(q, k, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, bpos,
+                             C.rms_eps)
+            att = torch.empty_like(q)
+            o.dflash_block_attn(att, q, k, v, self.ring_k[i], self.ring_v[i],
+                                self.ring_pos, bpos, scale)
+            x = x + (att.reshape(B, -1) @ self.w[f"wo.{i}"]).float()
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps).to(self.wdt)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + ((gate * (f @ self.w[f"ffn_up.{i}"]))
+                     @ self.w[f"ffn_down.{i}"]).float()
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return (x.to(self.wdt) @ self.w["lm_head"]).float()
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        return self.block_step(id_last, n_past, n_draft)[1:].argmax(dim=-1)
+    # ---- graphed step ---------------------------------------------------
+    def _static_buffers(self, B):
+        """Every tensor a captured step touches, allocated once and held."""
+        C = self.cfg
+        d = self.device
+        wdt = self.wdt
+        s = {
+            "tokens": torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                                 device=d),
+            "arange": torch.arange(B, dtype=torch.int32, device=d),
+            "n_past": torch.zeros((), dtype=torch.int32, device=d),
+            "bpos": torch.empty(B, dtype=torch.int32, device=d),
+            "xw": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "x": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "sq": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "ms": torch.empty(B, 1, dtype=torch.float32, device=d),
+            "hn": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "h": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "q": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "k": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "v": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "att": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "ao": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "g1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "sg": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "u1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "d1": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "lg": torch.empty(B, C.n_vocab, dtype=wdt, device=d),
+            "logits": torch.empty(B, C.n_vocab, dtype=torch.float32,
+                                  device=d),
+        }
+        return s
+    def _rms_static(self, src, w, s, out):
+        """RMS norm of `src` into `out` using only held buffers."""
+        torch.pow(src, 2, out=s["sq"])
+        torch.mean(s["sq"], -1, keepdim=True, out=s["ms"])
+        s["ms"].add_(self.cfg.rms_eps).rsqrt_()
+        torch.mul(src, s["ms"], out=out)
+        out.mul_(w)
+    def _forward_static(self, s):
+        """Allocation-free block forward into held buffers; numerically
+        identical to `_forward`."""
+        C = self.cfg
+        o = _ops()
+        B = s["tokens"].numel()
+        scale = C.head_dim ** -0.5
+        torch.index_select(self.w["tok_embd"], 0, s["tokens"], out=s["xw"])
+        s["x"].copy_(s["xw"])
+        for i in range(C.n_layer):
+            self._rms_static(s["x"], self.w[f"attn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"wq.{i}"], out=s["q"])
+            torch.matmul(s["h"], self.w[f"wk.{i}"], out=s["k"])
+            torch.matmul(s["h"], self.w[f"wv.{i}"], out=s["v"])
+            qv = s["q"].view(B, C.n_head, C.head_dim)
+            kv_ = s["k"].view(B, C.n_head_kv, C.head_dim)
+            vv = s["v"].view(B, C.n_head_kv, C.head_dim)
+            o.dflash_qk_rope(qv, kv_, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, s["bpos"],
+                             C.rms_eps)
+            av = s["att"].view(B, C.n_head, C.head_dim)
+            o.dflash_block_attn(av, qv, kv_, vv, self.ring_k[i],
+                                self.ring_v[i], self.ring_pos, s["bpos"],
+                                scale)
+            torch.matmul(s["att"], self.w[f"wo.{i}"], out=s["ao"])
+            s["x"].add_(s["ao"])
+            self._rms_static(s["x"], self.w[f"ffn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"ffn_gate.{i}"], out=s["g1"])
+            # fused silu keeps the step bitwise-equal to the eager path
+            torch.ops.aten.silu.out(s["g1"], out=s["sg"])
+            torch.matmul(s["h"], self.w[f"ffn_up.{i}"], out=s["u1"])
+            s["sg"].mul_(s["u1"])
+            torch.matmul(s["sg"], self.w[f"ffn_down.{i}"], out=s["d1"])
+            s["x"].add_(s["d1"])
+        self._rms_static(s["x"], self.w["out_norm"], s, s["hn"])
+        s["h"].copy_(s["hn"])
+        torch.matmul(s["h"], self.w["lm_head"], out=s["lg"])
+        s["logits"].copy_(s["lg"])
+        return s["logits"]
+    def capture_step(self, n_draft=None):
+        """Graph-captured block step over static buffers (the region allocates
+        nothing); returns step(id_last, n_past) -> logits, overwritten per replay."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        s = self._static_buffers(n_draft + 1)
+        stream = torch.cuda.Stream(device=self.device)
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):                    # warm kernels and cuBLAS
+                torch.add(s["arange"], s["n_past"], out=s["bpos"])
+                self._forward_static(s)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            torch.add(s["arange"], s["n_past"], out=s["bpos"])  # in-graph: replay sees updates
+            logits = self._forward_static(s)
+        def step(id_last, n_past):
+            s["tokens"][0].fill_(int(id_last))
+            s["n_past"].fill_(int(n_past))
+            graph.replay()
+            return logits
+        step.graph = graph
+        step.buffers = s
+        return step
+    # ---- verification ---------------------------------------------------
+    def accept(self, draft, target):
+        """Leading agreement count, device-side."""
+        o = _ops()
+        n = torch.zeros(1, dtype=torch.int32, device=self.device)
+        o.dflash_accept(n, draft.to(torch.int32).contiguous(),
+                        target.to(torch.int32).contiguous())
+        return n

build/torch212-cxx11-cu132-x86_64-linux/gguf_io.py ADDED Viewed

	@@ -0,0 +1,61 @@

+"""Load the DFlash drafter GGUF into reference-layout fp32 tensors; 2-D weights
+transpose from GGUF [out, in] to [in, out]; tok_embd and lm_head are the
+caller's (they belong to the target model)."""
+import numpy as np
+import torch
+from gguf import GGUFReader
+from gguf import quants as gq
+from .reference import DFlashConfig
+_BLK = {
+    "attn_norm": "attn_norm", "attn_q": "wq", "attn_k": "wk", "attn_v": "wv",
+    "attn_output": "wo", "attn_q_norm": "q_norm", "attn_k_norm": "k_norm",
+    "ffn_norm": "ffn_norm", "ffn_gate": "ffn_gate", "ffn_up": "ffn_up",
+    "ffn_down": "ffn_down",
+}
+_TOP = {"fc": "fc", "enc.output_norm": "enc_norm", "output_norm": "out_norm"}
+def _field(reader, key, default=None):
+    f = reader.get_field(key)
+    if f is None:
+        return default
+    v = f.contents()
+    return v
+def load_drafter(path):
+    """-> (DFlashConfig, weights dict) with fp32 tensors in reference layout."""
+    r = GGUFReader(str(path))
+    cfg = DFlashConfig(
+        n_layer=int(_field(r, "dflash.block_count")),
+        n_embd=int(_field(r, "dflash.embedding_length")),
+        n_head=int(_field(r, "dflash.attention.head_count")),
+        n_head_kv=int(_field(r, "dflash.attention.head_count_kv")),
+        head_dim=int(_field(r, "dflash.attention.key_length")),
+        n_ff=int(_field(r, "dflash.feed_forward_length")),
+        rope_base=float(_field(r, "dflash.rope.freq_base")),
+        rms_eps=float(_field(r, "dflash.attention.layer_norm_rms_epsilon")),
+        window=int(_field(r, "dflash.attention.sliding_window")),
+        block_size=int(_field(r, "dflash.block_size")),
+        target_layers=tuple(_field(r, "dflash.target_layers")),
+    )
+    w = {}
+    for t in r.tensors:
+        name = t.name
+        assert name.endswith(".weight"), name
+        base = name[: -len(".weight")]
+        if base.startswith("blk."):
+            _, idx, part = base.split(".", 2)
+            key = f"{_BLK[part]}.{idx}"
+        else:
+            key = _TOP[base]
+        data = gq.dequantize(t.data, t.tensor_type)
+        arr = torch.from_numpy(np.ascontiguousarray(data)).float()
+        if arr.dim() == 2:
+            arr = arr.t().contiguous()
+        w[key] = arr
+    expect = 3 + 11 * cfg.n_layer
+    assert len(w) == expect, (len(w), expect)
+    return cfg, w

build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "dflash-ops",
+  "id": "_dflash_ops_cuda_03c9586",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "wNpuIbQEXTfFNjrmz6aISM46VFgMKAHkqqu+ECJJS20=",
+      "_dflash_ops_cuda_03c9586.abi3.so": "zZG+F2T7Md0lecfy6fB5Inhz3lQFoWKxQyStSzXMHzU=",
+      "_ops.py": "+6e3Oo6lwRkZaeOrhSSZEwWYs4/Fpd/zQgJj4MVEi9I=",
+      "drafter.py": "wGU5TWlXWr/ZogxQr4pho5Wku66UUygnczpsTu82dls=",
+      "gguf_io.py": "XrgF2dH7LNR1KTNKKu/pKhh6wT90bCjcjCAgmyGsuZA=",
+      "reference.py": "InmFS9rUwURZ7ZGVUxwQJCCjJZR/hKWrF3eMfp4zQvY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "03c9586154cc2afa6efaa7b111dccef715ec9d69",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/reference.py ADDED Viewed

	@@ -0,0 +1,188 @@

+"""Eager fp32 DFlash drafter reference, transcribed from llama.cpp's
+llama_model_dflash graph; the CUDA ops are certified against this module."""
+from dataclasses import dataclass, field
+import torch
+@dataclass
+class DFlashConfig:
+    n_layer: int = 5
+    n_embd: int = 6656
+    n_head: int = 32
+    n_head_kv: int = 8
+    head_dim: int = 128
+    n_ff: int = 19968
+    n_vocab: int = 202048
+    rope_base: float = 500000.0
+    rms_eps: float = 1e-6
+    window: int = 2048            # SWA width; also the ring capacity
+    block_size: int = 16
+    mask_token_id: int = 201818
+    target_layers: tuple = (2, 14, 26, 38, 50)
+    @property
+    def n_feat(self):
+        return len(self.target_layers) * self.n_embd
+    @classmethod
+    def tiny(cls):
+        """Toy sizes with the real structure; head_dim stays a multiple of 32."""
+        return cls(n_layer=2, n_embd=64, n_head=4, n_head_kv=2, head_dim=32,
+                   n_ff=96, n_vocab=997, window=32, block_size=8,
+                   mask_token_id=990, target_layers=(1, 2))
+def rms_norm(x, w, eps):
+    """llama.cpp LLM_NORM_RMS: x * w / sqrt(mean(x^2) + eps), fp32."""
+    v = x.float()
+    return (v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps)) * w.float()
+def rope_inv_freq(head_dim, base, device="cpu"):
+    """fp64 NEOX inverse frequencies base^(-2i/head_dim), host-evaluated so the
+    table is device-invariant."""
+    half = head_dim // 2
+    t = base ** (-torch.arange(half, dtype=torch.float64) * 2.0 / head_dim)
+    return t.to(device)
+def rope_neox(x, pos, base=None, inv=None):
+    """NEOX rope over pairs (i, i + d/2), fp32; x [T, H, D], pos [T]; pass
+    `inv` to reuse a precomputed table."""
+    T, H, D = x.shape
+    half = D // 2
+    if inv is None:
+        inv = rope_inv_freq(D, base, x.device)
+    ang = pos.to(torch.float64)[:, None] * inv[None, :]          # [T, half]
+    cos = ang.cos().to(torch.float32)[:, None, :]
+    sin = ang.sin().to(torch.float32)[:, None, :]
+    xf = x.float()
+    x1, x2 = xf[..., :half], xf[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+class DFlashRef:
+    """Reference drafter over an fp32 weight dict (layout as in
+    `random_weights`) with a sliding ring KV state."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cpu",
+                 ring_dtype="fp32"):
+        self.cfg = cfg
+        self.w = {k: v.to(device=device, dtype=torch.float32)
+                  for k, v in weights.items()}
+        self.device = torch.device(device)
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base, device)
+        # per-layer ring K/V [window, Hkv, D]; ring_pos = slot's position, -1 empty
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int64,
+                                   device=device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        """Fuse per-position target features [T, n_feat] to g [T, n_embd]."""
+        g = feats.float() @ self.w["fc"]
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Project fused features into every layer's ring at `pos` [T], chunked
+        at `window` so no two tokens in one pass share a slot."""
+        C = self.cfg
+        for beg in range(0, len(pos), C.window):
+            p = pos[beg:beg + C.window]
+            gc = g[beg:beg + C.window]
+            slots = p % C.window
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+                k = rope_neox(k, p, inv=self.rope_inv)
+                self.ring_k[i, slots] = k.to(self.ring_dt)
+                self.ring_v[i, slots] = v.to(self.ring_dt)
+            self.ring_pos[slots] = p
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        """Denoise [id_last, mask x n_draft] at n_past..; -> [n_draft+1, n_vocab]."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = id_last
+        pos = n_past + torch.arange(B, device=self.device)
+        x = self.w["tok_embd"][tokens]
+        live = self.ring_pos >= 0
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head, C.head_dim)
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            q = rms_norm(q, self.w[f"q_norm.{i}"], C.rms_eps)
+            k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+            q = rope_neox(q, pos, inv=self.rope_inv)
+            k = rope_neox(k, pos, inv=self.rope_inv)
+            # keys = live ring + block; non-causal, SWA masks only the deep past
+            keys = torch.cat([self.ring_k[i, live].float(), k], dim=0)
+            vals = torch.cat([self.ring_v[i, live].float(), v], dim=0)
+            kpos = torch.cat([self.ring_pos[live], pos], dim=0)
+            gq = C.n_head // C.n_head_kv
+            qh = q.permute(1, 0, 2)                                # [H, B, D]
+            kh = keys.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            vh = vals.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            att = qh @ kh.transpose(1, 2) / (C.head_dim ** 0.5)    # [H, B, S]
+            masked = (pos[:, None] - kpos[None, :]) >= C.window    # SWA
+            att = att.masked_fill(masked[None, :, :], float("-inf"))
+            att = att.softmax(dim=-1)
+            out = (att @ vh).permute(1, 0, 2).reshape(B, -1)       # [B, H*D]
+            x = x + out @ self.w[f"wo.{i}"]
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + (gate * (f @ self.w[f"ffn_up.{i}"])) @ self.w[f"ffn_down.{i}"]
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return x @ self.w["lm_head"]
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        """Greedy drafts from block positions 1..n_draft (position 0 unused)."""
+        logits = self.block_step(id_last, n_past, n_draft)
+        return logits[1:].argmax(dim=-1)
+def random_weights(cfg: DFlashConfig, seed=0, device="cpu"):
+    """Gaussian weights at transformer-typical scale, for structural tests."""
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    def t(*shape, s=0.02):
+        return torch.randn(*shape, generator=g) * s
+    C = cfg
+    w = {"fc": t(C.n_feat, C.n_embd), "enc_norm": torch.ones(C.n_embd),
+         "out_norm": torch.ones(C.n_embd),
+         "tok_embd": t(C.n_vocab, C.n_embd),
+         "lm_head": t(C.n_embd, C.n_vocab)}
+    for i in range(C.n_layer):
+        w[f"attn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"wq.{i}"] = t(C.n_embd, C.n_head * C.head_dim)
+        w[f"wk.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wv.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wo.{i}"] = t(C.n_head * C.head_dim, C.n_embd)
+        w[f"q_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"k_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"ffn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"ffn_gate.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_up.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_down.{i}"] = t(C.n_ff, C.n_embd)
+    return {k: v.to(device) for k, v in w.items()}

build/torch213-cxx11-cu126-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""DFlash block-diffusion drafter ops; the eager reference defines the
+semantics and the CUDA ops are certified against it."""
+from .reference import DFlashConfig, DFlashRef, random_weights  # noqa: F401
+from .drafter import DFlashDrafter  # noqa: F401
+try:  # built variants ship _ops
+    from ._ops import ops  # noqa: F401
+except ImportError:
+    try:  # a local JIT build registers the same ops under torch.ops.dflash_ext
+        import torch as _torch
+        _torch.ops.dflash_ext.dflash_accept
+        ops = _torch.ops.dflash_ext
+    except Exception:
+        ops = None
+try:
+    from .gguf_io import load_drafter  # noqa: F401
+except ImportError:  # gguf-py absent; the reference stays usable
+    load_drafter = None
+__all__ = ["DFlashConfig", "DFlashRef", "DFlashDrafter", "random_weights",
+           "load_drafter", "ops"]

build/torch213-cxx11-cu126-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:9584d73d21d10eb16f12f42a0e5658fa8c2718b4df0dd748ba5b6c6e0f41b7ac
+size 893256

build/torch213-cxx11-cu126-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _dflash_ops_cuda_03c9586
+ops = torch.ops._dflash_ops_cuda_03c9586
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_dflash_ops_cuda_03c9586::{op_name}"

build/torch213-cxx11-cu126-x86_64-linux/drafter.py ADDED Viewed

	@@ -0,0 +1,236 @@

+"""GPU DFlash drafter: cuBLAS GEMMs plus the fused ops, with a CUDA-graphed
+block step."""
+import torch
+from .reference import DFlashConfig, rms_norm, rope_inv_freq
+def _ops():
+    from . import ops
+    if ops is None:
+        raise RuntimeError(
+            "dflash-ops compiled kernels are unavailable; build them with "
+            "load_local.load() in a source checkout, or install a published "
+            "variant with kernels.get_kernel")
+    return ops
+class DFlashDrafter:
+    """Ring-cached DFlash drafter on CUDA. Weight keys match `DFlashRef`."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cuda",
+                 ring_dtype="fp32", dtype="fp32"):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        # dtype sets GEMM weights and activations; residual, norms, softmax stay fp32
+        self.wdt = torch.bfloat16 if dtype == "bf16" else torch.float32
+        self.w = {}
+        for k, v in weights.items():
+            dt = self.wdt if v.dim() == 2 else torch.float32
+            self.w[k] = v.to(device=self.device, dtype=dt).contiguous()
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        # host-built fp64 rope table, shared by every op call
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base,
+                                      self.device).contiguous()
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=self.device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int32,
+                                   device=self.device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        g = (feats.to(self.device, self.wdt) @ self.w["fc"]).float()
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Ring injection at `pos`, chunked at `window` so no slot is written twice per launch."""
+        C = self.cfg
+        o = _ops()
+        pos = pos.to(self.device, torch.int32).contiguous()
+        for beg in range(0, pos.numel(), C.window):
+            p = pos[beg:beg + C.window].contiguous()
+            gc = g[beg:beg + C.window].to(self.wdt)
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                o.dflash_inject(self.ring_k[i], self.ring_v[i], self.ring_pos,
+                                k, v, self.w[f"k_norm.{i}"], self.rope_inv, p,
+                                C.rms_eps)
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = int(id_last)
+        bpos = (n_past + torch.arange(B, device=self.device)).to(torch.int32)
+        return self._forward(tokens, bpos.contiguous())
+    def _forward(self, tokens, bpos):
+        """Block forward on caller-owned device tensors; reads no host state."""
+        C = self.cfg
+        o = _ops()
+        B = tokens.numel()
+        x = self.w["tok_embd"][tokens].float()
+        scale = C.head_dim ** -0.5
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps).to(self.wdt)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head,
+                                             C.head_dim).contiguous()
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            o.dflash_qk_rope(q, k, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, bpos,
+                             C.rms_eps)
+            att = torch.empty_like(q)
+            o.dflash_block_attn(att, q, k, v, self.ring_k[i], self.ring_v[i],
+                                self.ring_pos, bpos, scale)
+            x = x + (att.reshape(B, -1) @ self.w[f"wo.{i}"]).float()
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps).to(self.wdt)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + ((gate * (f @ self.w[f"ffn_up.{i}"]))
+                     @ self.w[f"ffn_down.{i}"]).float()
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return (x.to(self.wdt) @ self.w["lm_head"]).float()
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        return self.block_step(id_last, n_past, n_draft)[1:].argmax(dim=-1)
+    # ---- graphed step ---------------------------------------------------
+    def _static_buffers(self, B):
+        """Every tensor a captured step touches, allocated once and held."""
+        C = self.cfg
+        d = self.device
+        wdt = self.wdt
+        s = {
+            "tokens": torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                                 device=d),
+            "arange": torch.arange(B, dtype=torch.int32, device=d),
+            "n_past": torch.zeros((), dtype=torch.int32, device=d),
+            "bpos": torch.empty(B, dtype=torch.int32, device=d),
+            "xw": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "x": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "sq": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "ms": torch.empty(B, 1, dtype=torch.float32, device=d),
+            "hn": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "h": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "q": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "k": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "v": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "att": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "ao": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "g1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "sg": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "u1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "d1": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "lg": torch.empty(B, C.n_vocab, dtype=wdt, device=d),
+            "logits": torch.empty(B, C.n_vocab, dtype=torch.float32,
+                                  device=d),
+        }
+        return s
+    def _rms_static(self, src, w, s, out):
+        """RMS norm of `src` into `out` using only held buffers."""
+        torch.pow(src, 2, out=s["sq"])
+        torch.mean(s["sq"], -1, keepdim=True, out=s["ms"])
+        s["ms"].add_(self.cfg.rms_eps).rsqrt_()
+        torch.mul(src, s["ms"], out=out)
+        out.mul_(w)
+    def _forward_static(self, s):
+        """Allocation-free block forward into held buffers; numerically
+        identical to `_forward`."""
+        C = self.cfg
+        o = _ops()
+        B = s["tokens"].numel()
+        scale = C.head_dim ** -0.5
+        torch.index_select(self.w["tok_embd"], 0, s["tokens"], out=s["xw"])
+        s["x"].copy_(s["xw"])
+        for i in range(C.n_layer):
+            self._rms_static(s["x"], self.w[f"attn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"wq.{i}"], out=s["q"])
+            torch.matmul(s["h"], self.w[f"wk.{i}"], out=s["k"])
+            torch.matmul(s["h"], self.w[f"wv.{i}"], out=s["v"])
+            qv = s["q"].view(B, C.n_head, C.head_dim)
+            kv_ = s["k"].view(B, C.n_head_kv, C.head_dim)
+            vv = s["v"].view(B, C.n_head_kv, C.head_dim)
+            o.dflash_qk_rope(qv, kv_, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, s["bpos"],
+                             C.rms_eps)
+            av = s["att"].view(B, C.n_head, C.head_dim)
+            o.dflash_block_attn(av, qv, kv_, vv, self.ring_k[i],
+                                self.ring_v[i], self.ring_pos, s["bpos"],
+                                scale)
+            torch.matmul(s["att"], self.w[f"wo.{i}"], out=s["ao"])
+            s["x"].add_(s["ao"])
+            self._rms_static(s["x"], self.w[f"ffn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"ffn_gate.{i}"], out=s["g1"])
+            # fused silu keeps the step bitwise-equal to the eager path
+            torch.ops.aten.silu.out(s["g1"], out=s["sg"])
+            torch.matmul(s["h"], self.w[f"ffn_up.{i}"], out=s["u1"])
+            s["sg"].mul_(s["u1"])
+            torch.matmul(s["sg"], self.w[f"ffn_down.{i}"], out=s["d1"])
+            s["x"].add_(s["d1"])
+        self._rms_static(s["x"], self.w["out_norm"], s, s["hn"])
+        s["h"].copy_(s["hn"])
+        torch.matmul(s["h"], self.w["lm_head"], out=s["lg"])
+        s["logits"].copy_(s["lg"])
+        return s["logits"]
+    def capture_step(self, n_draft=None):
+        """Graph-captured block step over static buffers (the region allocates
+        nothing); returns step(id_last, n_past) -> logits, overwritten per replay."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        s = self._static_buffers(n_draft + 1)
+        stream = torch.cuda.Stream(device=self.device)
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):                    # warm kernels and cuBLAS
+                torch.add(s["arange"], s["n_past"], out=s["bpos"])
+                self._forward_static(s)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            torch.add(s["arange"], s["n_past"], out=s["bpos"])  # in-graph: replay sees updates
+            logits = self._forward_static(s)
+        def step(id_last, n_past):
+            s["tokens"][0].fill_(int(id_last))
+            s["n_past"].fill_(int(n_past))
+            graph.replay()
+            return logits
+        step.graph = graph
+        step.buffers = s
+        return step
+    # ---- verification ---------------------------------------------------
+    def accept(self, draft, target):
+        """Leading agreement count, device-side."""
+        o = _ops()
+        n = torch.zeros(1, dtype=torch.int32, device=self.device)
+        o.dflash_accept(n, draft.to(torch.int32).contiguous(),
+                        target.to(torch.int32).contiguous())
+        return n

build/torch213-cxx11-cu126-x86_64-linux/gguf_io.py ADDED Viewed

	@@ -0,0 +1,61 @@

+"""Load the DFlash drafter GGUF into reference-layout fp32 tensors; 2-D weights
+transpose from GGUF [out, in] to [in, out]; tok_embd and lm_head are the
+caller's (they belong to the target model)."""
+import numpy as np
+import torch
+from gguf import GGUFReader
+from gguf import quants as gq
+from .reference import DFlashConfig
+_BLK = {
+    "attn_norm": "attn_norm", "attn_q": "wq", "attn_k": "wk", "attn_v": "wv",
+    "attn_output": "wo", "attn_q_norm": "q_norm", "attn_k_norm": "k_norm",
+    "ffn_norm": "ffn_norm", "ffn_gate": "ffn_gate", "ffn_up": "ffn_up",
+    "ffn_down": "ffn_down",
+}
+_TOP = {"fc": "fc", "enc.output_norm": "enc_norm", "output_norm": "out_norm"}
+def _field(reader, key, default=None):
+    f = reader.get_field(key)
+    if f is None:
+        return default
+    v = f.contents()
+    return v
+def load_drafter(path):
+    """-> (DFlashConfig, weights dict) with fp32 tensors in reference layout."""
+    r = GGUFReader(str(path))
+    cfg = DFlashConfig(
+        n_layer=int(_field(r, "dflash.block_count")),
+        n_embd=int(_field(r, "dflash.embedding_length")),
+        n_head=int(_field(r, "dflash.attention.head_count")),
+        n_head_kv=int(_field(r, "dflash.attention.head_count_kv")),
+        head_dim=int(_field(r, "dflash.attention.key_length")),
+        n_ff=int(_field(r, "dflash.feed_forward_length")),
+        rope_base=float(_field(r, "dflash.rope.freq_base")),
+        rms_eps=float(_field(r, "dflash.attention.layer_norm_rms_epsilon")),
+        window=int(_field(r, "dflash.attention.sliding_window")),
+        block_size=int(_field(r, "dflash.block_size")),
+        target_layers=tuple(_field(r, "dflash.target_layers")),
+    )
+    w = {}
+    for t in r.tensors:
+        name = t.name
+        assert name.endswith(".weight"), name
+        base = name[: -len(".weight")]
+        if base.startswith("blk."):
+            _, idx, part = base.split(".", 2)
+            key = f"{_BLK[part]}.{idx}"
+        else:
+            key = _TOP[base]
+        data = gq.dequantize(t.data, t.tensor_type)
+        arr = torch.from_numpy(np.ascontiguousarray(data)).float()
+        if arr.dim() == 2:
+            arr = arr.t().contiguous()
+        w[key] = arr
+    expect = 3 + 11 * cfg.n_layer
+    assert len(w) == expect, (len(w), expect)
+    return cfg, w

build/torch213-cxx11-cu126-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "dflash-ops",
+  "id": "_dflash_ops_cuda_03c9586",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "wNpuIbQEXTfFNjrmz6aISM46VFgMKAHkqqu+ECJJS20=",
+      "_dflash_ops_cuda_03c9586.abi3.so": "lYTXPSHRDrFvEvQqDlZY+ownGLTfDddIultsbg9Bt6w=",
+      "_ops.py": "+6e3Oo6lwRkZaeOrhSSZEwWYs4/Fpd/zQgJj4MVEi9I=",
+      "drafter.py": "wGU5TWlXWr/ZogxQr4pho5Wku66UUygnczpsTu82dls=",
+      "gguf_io.py": "XrgF2dH7LNR1KTNKKu/pKhh6wT90bCjcjCAgmyGsuZA=",
+      "reference.py": "InmFS9rUwURZ7ZGVUxwQJCCjJZR/hKWrF3eMfp4zQvY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "03c9586154cc2afa6efaa7b111dccef715ec9d69",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu126-x86_64-linux/reference.py ADDED Viewed

	@@ -0,0 +1,188 @@

+"""Eager fp32 DFlash drafter reference, transcribed from llama.cpp's
+llama_model_dflash graph; the CUDA ops are certified against this module."""
+from dataclasses import dataclass, field
+import torch
+@dataclass
+class DFlashConfig:
+    n_layer: int = 5
+    n_embd: int = 6656
+    n_head: int = 32
+    n_head_kv: int = 8
+    head_dim: int = 128
+    n_ff: int = 19968
+    n_vocab: int = 202048
+    rope_base: float = 500000.0
+    rms_eps: float = 1e-6
+    window: int = 2048            # SWA width; also the ring capacity
+    block_size: int = 16
+    mask_token_id: int = 201818
+    target_layers: tuple = (2, 14, 26, 38, 50)
+    @property
+    def n_feat(self):
+        return len(self.target_layers) * self.n_embd
+    @classmethod
+    def tiny(cls):
+        """Toy sizes with the real structure; head_dim stays a multiple of 32."""
+        return cls(n_layer=2, n_embd=64, n_head=4, n_head_kv=2, head_dim=32,
+                   n_ff=96, n_vocab=997, window=32, block_size=8,
+                   mask_token_id=990, target_layers=(1, 2))
+def rms_norm(x, w, eps):
+    """llama.cpp LLM_NORM_RMS: x * w / sqrt(mean(x^2) + eps), fp32."""
+    v = x.float()
+    return (v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps)) * w.float()
+def rope_inv_freq(head_dim, base, device="cpu"):
+    """fp64 NEOX inverse frequencies base^(-2i/head_dim), host-evaluated so the
+    table is device-invariant."""
+    half = head_dim // 2
+    t = base ** (-torch.arange(half, dtype=torch.float64) * 2.0 / head_dim)
+    return t.to(device)
+def rope_neox(x, pos, base=None, inv=None):
+    """NEOX rope over pairs (i, i + d/2), fp32; x [T, H, D], pos [T]; pass
+    `inv` to reuse a precomputed table."""
+    T, H, D = x.shape
+    half = D // 2
+    if inv is None:
+        inv = rope_inv_freq(D, base, x.device)
+    ang = pos.to(torch.float64)[:, None] * inv[None, :]          # [T, half]
+    cos = ang.cos().to(torch.float32)[:, None, :]
+    sin = ang.sin().to(torch.float32)[:, None, :]
+    xf = x.float()
+    x1, x2 = xf[..., :half], xf[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+class DFlashRef:
+    """Reference drafter over an fp32 weight dict (layout as in
+    `random_weights`) with a sliding ring KV state."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cpu",
+                 ring_dtype="fp32"):
+        self.cfg = cfg
+        self.w = {k: v.to(device=device, dtype=torch.float32)
+                  for k, v in weights.items()}
+        self.device = torch.device(device)
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base, device)
+        # per-layer ring K/V [window, Hkv, D]; ring_pos = slot's position, -1 empty
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int64,
+                                   device=device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        """Fuse per-position target features [T, n_feat] to g [T, n_embd]."""
+        g = feats.float() @ self.w["fc"]
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Project fused features into every layer's ring at `pos` [T], chunked
+        at `window` so no two tokens in one pass share a slot."""
+        C = self.cfg
+        for beg in range(0, len(pos), C.window):
+            p = pos[beg:beg + C.window]
+            gc = g[beg:beg + C.window]
+            slots = p % C.window
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+                k = rope_neox(k, p, inv=self.rope_inv)
+                self.ring_k[i, slots] = k.to(self.ring_dt)
+                self.ring_v[i, slots] = v.to(self.ring_dt)
+            self.ring_pos[slots] = p
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        """Denoise [id_last, mask x n_draft] at n_past..; -> [n_draft+1, n_vocab]."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = id_last
+        pos = n_past + torch.arange(B, device=self.device)
+        x = self.w["tok_embd"][tokens]
+        live = self.ring_pos >= 0
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head, C.head_dim)
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            q = rms_norm(q, self.w[f"q_norm.{i}"], C.rms_eps)
+            k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+            q = rope_neox(q, pos, inv=self.rope_inv)
+            k = rope_neox(k, pos, inv=self.rope_inv)
+            # keys = live ring + block; non-causal, SWA masks only the deep past
+            keys = torch.cat([self.ring_k[i, live].float(), k], dim=0)
+            vals = torch.cat([self.ring_v[i, live].float(), v], dim=0)
+            kpos = torch.cat([self.ring_pos[live], pos], dim=0)
+            gq = C.n_head // C.n_head_kv
+            qh = q.permute(1, 0, 2)                                # [H, B, D]
+            kh = keys.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            vh = vals.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            att = qh @ kh.transpose(1, 2) / (C.head_dim ** 0.5)    # [H, B, S]
+            masked = (pos[:, None] - kpos[None, :]) >= C.window    # SWA
+            att = att.masked_fill(masked[None, :, :], float("-inf"))
+            att = att.softmax(dim=-1)
+            out = (att @ vh).permute(1, 0, 2).reshape(B, -1)       # [B, H*D]
+            x = x + out @ self.w[f"wo.{i}"]
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + (gate * (f @ self.w[f"ffn_up.{i}"])) @ self.w[f"ffn_down.{i}"]
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return x @ self.w["lm_head"]
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        """Greedy drafts from block positions 1..n_draft (position 0 unused)."""
+        logits = self.block_step(id_last, n_past, n_draft)
+        return logits[1:].argmax(dim=-1)
+def random_weights(cfg: DFlashConfig, seed=0, device="cpu"):
+    """Gaussian weights at transformer-typical scale, for structural tests."""
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    def t(*shape, s=0.02):
+        return torch.randn(*shape, generator=g) * s
+    C = cfg
+    w = {"fc": t(C.n_feat, C.n_embd), "enc_norm": torch.ones(C.n_embd),
+         "out_norm": torch.ones(C.n_embd),
+         "tok_embd": t(C.n_vocab, C.n_embd),
+         "lm_head": t(C.n_embd, C.n_vocab)}
+    for i in range(C.n_layer):
+        w[f"attn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"wq.{i}"] = t(C.n_embd, C.n_head * C.head_dim)
+        w[f"wk.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wv.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wo.{i}"] = t(C.n_head * C.head_dim, C.n_embd)
+        w[f"q_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"k_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"ffn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"ffn_gate.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_up.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_down.{i}"] = t(C.n_ff, C.n_embd)
+    return {k: v.to(device) for k, v in w.items()}

build/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""DFlash block-diffusion drafter ops; the eager reference defines the
+semantics and the CUDA ops are certified against it."""
+from .reference import DFlashConfig, DFlashRef, random_weights  # noqa: F401
+from .drafter import DFlashDrafter  # noqa: F401
+try:  # built variants ship _ops
+    from ._ops import ops  # noqa: F401
+except ImportError:
+    try:  # a local JIT build registers the same ops under torch.ops.dflash_ext
+        import torch as _torch
+        _torch.ops.dflash_ext.dflash_accept
+        ops = _torch.ops.dflash_ext
+    except Exception:
+        ops = None
+try:
+    from .gguf_io import load_drafter  # noqa: F401
+except ImportError:  # gguf-py absent; the reference stays usable
+    load_drafter = None
+__all__ = ["DFlashConfig", "DFlashRef", "DFlashDrafter", "random_weights",
+           "load_drafter", "ops"]

build/torch213-cxx11-cu130-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:b084949c234bf47f45c68cc117ba53eeb7b6fe87854b8fde895d0598dc835641
+size 1402096

build/torch213-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _dflash_ops_cuda_03c9586
+ops = torch.ops._dflash_ops_cuda_03c9586
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_dflash_ops_cuda_03c9586::{op_name}"

build/torch213-cxx11-cu130-x86_64-linux/drafter.py ADDED Viewed

	@@ -0,0 +1,236 @@

+"""GPU DFlash drafter: cuBLAS GEMMs plus the fused ops, with a CUDA-graphed
+block step."""
+import torch
+from .reference import DFlashConfig, rms_norm, rope_inv_freq
+def _ops():
+    from . import ops
+    if ops is None:
+        raise RuntimeError(
+            "dflash-ops compiled kernels are unavailable; build them with "
+            "load_local.load() in a source checkout, or install a published "
+            "variant with kernels.get_kernel")
+    return ops
+class DFlashDrafter:
+    """Ring-cached DFlash drafter on CUDA. Weight keys match `DFlashRef`."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cuda",
+                 ring_dtype="fp32", dtype="fp32"):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        # dtype sets GEMM weights and activations; residual, norms, softmax stay fp32
+        self.wdt = torch.bfloat16 if dtype == "bf16" else torch.float32
+        self.w = {}
+        for k, v in weights.items():
+            dt = self.wdt if v.dim() == 2 else torch.float32
+            self.w[k] = v.to(device=self.device, dtype=dt).contiguous()
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        # host-built fp64 rope table, shared by every op call
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base,
+                                      self.device).contiguous()
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=self.device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int32,
+                                   device=self.device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        g = (feats.to(self.device, self.wdt) @ self.w["fc"]).float()
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Ring injection at `pos`, chunked at `window` so no slot is written twice per launch."""
+        C = self.cfg
+        o = _ops()
+        pos = pos.to(self.device, torch.int32).contiguous()
+        for beg in range(0, pos.numel(), C.window):
+            p = pos[beg:beg + C.window].contiguous()
+            gc = g[beg:beg + C.window].to(self.wdt)
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                o.dflash_inject(self.ring_k[i], self.ring_v[i], self.ring_pos,
+                                k, v, self.w[f"k_norm.{i}"], self.rope_inv, p,
+                                C.rms_eps)
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = int(id_last)
+        bpos = (n_past + torch.arange(B, device=self.device)).to(torch.int32)
+        return self._forward(tokens, bpos.contiguous())
+    def _forward(self, tokens, bpos):
+        """Block forward on caller-owned device tensors; reads no host state."""
+        C = self.cfg
+        o = _ops()
+        B = tokens.numel()
+        x = self.w["tok_embd"][tokens].float()
+        scale = C.head_dim ** -0.5
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps).to(self.wdt)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head,
+                                             C.head_dim).contiguous()
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            o.dflash_qk_rope(q, k, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, bpos,
+                             C.rms_eps)
+            att = torch.empty_like(q)
+            o.dflash_block_attn(att, q, k, v, self.ring_k[i], self.ring_v[i],
+                                self.ring_pos, bpos, scale)
+            x = x + (att.reshape(B, -1) @ self.w[f"wo.{i}"]).float()
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps).to(self.wdt)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + ((gate * (f @ self.w[f"ffn_up.{i}"]))
+                     @ self.w[f"ffn_down.{i}"]).float()
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return (x.to(self.wdt) @ self.w["lm_head"]).float()
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        return self.block_step(id_last, n_past, n_draft)[1:].argmax(dim=-1)
+    # ---- graphed step ---------------------------------------------------
+    def _static_buffers(self, B):
+        """Every tensor a captured step touches, allocated once and held."""
+        C = self.cfg
+        d = self.device
+        wdt = self.wdt
+        s = {
+            "tokens": torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                                 device=d),
+            "arange": torch.arange(B, dtype=torch.int32, device=d),
+            "n_past": torch.zeros((), dtype=torch.int32, device=d),
+            "bpos": torch.empty(B, dtype=torch.int32, device=d),
+            "xw": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "x": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "sq": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "ms": torch.empty(B, 1, dtype=torch.float32, device=d),
+            "hn": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "h": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "q": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "k": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "v": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "att": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "ao": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "g1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "sg": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "u1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "d1": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "lg": torch.empty(B, C.n_vocab, dtype=wdt, device=d),
+            "logits": torch.empty(B, C.n_vocab, dtype=torch.float32,
+                                  device=d),
+        }
+        return s
+    def _rms_static(self, src, w, s, out):
+        """RMS norm of `src` into `out` using only held buffers."""
+        torch.pow(src, 2, out=s["sq"])
+        torch.mean(s["sq"], -1, keepdim=True, out=s["ms"])
+        s["ms"].add_(self.cfg.rms_eps).rsqrt_()
+        torch.mul(src, s["ms"], out=out)
+        out.mul_(w)
+    def _forward_static(self, s):
+        """Allocation-free block forward into held buffers; numerically
+        identical to `_forward`."""
+        C = self.cfg
+        o = _ops()
+        B = s["tokens"].numel()
+        scale = C.head_dim ** -0.5
+        torch.index_select(self.w["tok_embd"], 0, s["tokens"], out=s["xw"])
+        s["x"].copy_(s["xw"])
+        for i in range(C.n_layer):
+            self._rms_static(s["x"], self.w[f"attn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"wq.{i}"], out=s["q"])
+            torch.matmul(s["h"], self.w[f"wk.{i}"], out=s["k"])
+            torch.matmul(s["h"], self.w[f"wv.{i}"], out=s["v"])
+            qv = s["q"].view(B, C.n_head, C.head_dim)
+            kv_ = s["k"].view(B, C.n_head_kv, C.head_dim)
+            vv = s["v"].view(B, C.n_head_kv, C.head_dim)
+            o.dflash_qk_rope(qv, kv_, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, s["bpos"],
+                             C.rms_eps)
+            av = s["att"].view(B, C.n_head, C.head_dim)
+            o.dflash_block_attn(av, qv, kv_, vv, self.ring_k[i],
+                                self.ring_v[i], self.ring_pos, s["bpos"],
+                                scale)
+            torch.matmul(s["att"], self.w[f"wo.{i}"], out=s["ao"])
+            s["x"].add_(s["ao"])
+            self._rms_static(s["x"], self.w[f"ffn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"ffn_gate.{i}"], out=s["g1"])
+            # fused silu keeps the step bitwise-equal to the eager path
+            torch.ops.aten.silu.out(s["g1"], out=s["sg"])
+            torch.matmul(s["h"], self.w[f"ffn_up.{i}"], out=s["u1"])
+            s["sg"].mul_(s["u1"])
+            torch.matmul(s["sg"], self.w[f"ffn_down.{i}"], out=s["d1"])
+            s["x"].add_(s["d1"])
+        self._rms_static(s["x"], self.w["out_norm"], s, s["hn"])
+        s["h"].copy_(s["hn"])
+        torch.matmul(s["h"], self.w["lm_head"], out=s["lg"])
+        s["logits"].copy_(s["lg"])
+        return s["logits"]
+    def capture_step(self, n_draft=None):
+        """Graph-captured block step over static buffers (the region allocates
+        nothing); returns step(id_last, n_past) -> logits, overwritten per replay."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        s = self._static_buffers(n_draft + 1)
+        stream = torch.cuda.Stream(device=self.device)
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):                    # warm kernels and cuBLAS
+                torch.add(s["arange"], s["n_past"], out=s["bpos"])
+                self._forward_static(s)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            torch.add(s["arange"], s["n_past"], out=s["bpos"])  # in-graph: replay sees updates
+            logits = self._forward_static(s)
+        def step(id_last, n_past):
+            s["tokens"][0].fill_(int(id_last))
+            s["n_past"].fill_(int(n_past))
+            graph.replay()
+            return logits
+        step.graph = graph
+        step.buffers = s
+        return step
+    # ---- verification ---------------------------------------------------
+    def accept(self, draft, target):
+        """Leading agreement count, device-side."""
+        o = _ops()
+        n = torch.zeros(1, dtype=torch.int32, device=self.device)
+        o.dflash_accept(n, draft.to(torch.int32).contiguous(),
+                        target.to(torch.int32).contiguous())
+        return n

build/torch213-cxx11-cu130-x86_64-linux/gguf_io.py ADDED Viewed

	@@ -0,0 +1,61 @@

+"""Load the DFlash drafter GGUF into reference-layout fp32 tensors; 2-D weights
+transpose from GGUF [out, in] to [in, out]; tok_embd and lm_head are the
+caller's (they belong to the target model)."""
+import numpy as np
+import torch
+from gguf import GGUFReader
+from gguf import quants as gq
+from .reference import DFlashConfig
+_BLK = {
+    "attn_norm": "attn_norm", "attn_q": "wq", "attn_k": "wk", "attn_v": "wv",
+    "attn_output": "wo", "attn_q_norm": "q_norm", "attn_k_norm": "k_norm",
+    "ffn_norm": "ffn_norm", "ffn_gate": "ffn_gate", "ffn_up": "ffn_up",
+    "ffn_down": "ffn_down",
+}
+_TOP = {"fc": "fc", "enc.output_norm": "enc_norm", "output_norm": "out_norm"}
+def _field(reader, key, default=None):
+    f = reader.get_field(key)
+    if f is None:
+        return default
+    v = f.contents()
+    return v
+def load_drafter(path):
+    """-> (DFlashConfig, weights dict) with fp32 tensors in reference layout."""
+    r = GGUFReader(str(path))
+    cfg = DFlashConfig(
+        n_layer=int(_field(r, "dflash.block_count")),
+        n_embd=int(_field(r, "dflash.embedding_length")),
+        n_head=int(_field(r, "dflash.attention.head_count")),
+        n_head_kv=int(_field(r, "dflash.attention.head_count_kv")),
+        head_dim=int(_field(r, "dflash.attention.key_length")),
+        n_ff=int(_field(r, "dflash.feed_forward_length")),
+        rope_base=float(_field(r, "dflash.rope.freq_base")),
+        rms_eps=float(_field(r, "dflash.attention.layer_norm_rms_epsilon")),
+        window=int(_field(r, "dflash.attention.sliding_window")),
+        block_size=int(_field(r, "dflash.block_size")),
+        target_layers=tuple(_field(r, "dflash.target_layers")),
+    )
+    w = {}
+    for t in r.tensors:
+        name = t.name
+        assert name.endswith(".weight"), name
+        base = name[: -len(".weight")]
+        if base.startswith("blk."):
+            _, idx, part = base.split(".", 2)
+            key = f"{_BLK[part]}.{idx}"
+        else:
+            key = _TOP[base]
+        data = gq.dequantize(t.data, t.tensor_type)
+        arr = torch.from_numpy(np.ascontiguousarray(data)).float()
+        if arr.dim() == 2:
+            arr = arr.t().contiguous()
+        w[key] = arr
+    expect = 3 + 11 * cfg.n_layer
+    assert len(w) == expect, (len(w), expect)
+    return cfg, w

build/torch213-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "dflash-ops",
+  "id": "_dflash_ops_cuda_03c9586",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "wNpuIbQEXTfFNjrmz6aISM46VFgMKAHkqqu+ECJJS20=",
+      "_dflash_ops_cuda_03c9586.abi3.so": "sISUnCNL9H9FxozBF7pT7re2/oeFS4/eiV0FmNyDVkE=",
+      "_ops.py": "+6e3Oo6lwRkZaeOrhSSZEwWYs4/Fpd/zQgJj4MVEi9I=",
+      "drafter.py": "wGU5TWlXWr/ZogxQr4pho5Wku66UUygnczpsTu82dls=",
+      "gguf_io.py": "XrgF2dH7LNR1KTNKKu/pKhh6wT90bCjcjCAgmyGsuZA=",
+      "reference.py": "InmFS9rUwURZ7ZGVUxwQJCCjJZR/hKWrF3eMfp4zQvY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "03c9586154cc2afa6efaa7b111dccef715ec9d69",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu130-x86_64-linux/reference.py ADDED Viewed

	@@ -0,0 +1,188 @@

+"""Eager fp32 DFlash drafter reference, transcribed from llama.cpp's
+llama_model_dflash graph; the CUDA ops are certified against this module."""
+from dataclasses import dataclass, field
+import torch
+@dataclass
+class DFlashConfig:
+    n_layer: int = 5
+    n_embd: int = 6656
+    n_head: int = 32
+    n_head_kv: int = 8
+    head_dim: int = 128
+    n_ff: int = 19968
+    n_vocab: int = 202048
+    rope_base: float = 500000.0
+    rms_eps: float = 1e-6
+    window: int = 2048            # SWA width; also the ring capacity
+    block_size: int = 16
+    mask_token_id: int = 201818
+    target_layers: tuple = (2, 14, 26, 38, 50)
+    @property
+    def n_feat(self):
+        return len(self.target_layers) * self.n_embd
+    @classmethod
+    def tiny(cls):
+        """Toy sizes with the real structure; head_dim stays a multiple of 32."""
+        return cls(n_layer=2, n_embd=64, n_head=4, n_head_kv=2, head_dim=32,
+                   n_ff=96, n_vocab=997, window=32, block_size=8,
+                   mask_token_id=990, target_layers=(1, 2))
+def rms_norm(x, w, eps):
+    """llama.cpp LLM_NORM_RMS: x * w / sqrt(mean(x^2) + eps), fp32."""
+    v = x.float()
+    return (v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps)) * w.float()
+def rope_inv_freq(head_dim, base, device="cpu"):
+    """fp64 NEOX inverse frequencies base^(-2i/head_dim), host-evaluated so the
+    table is device-invariant."""
+    half = head_dim // 2
+    t = base ** (-torch.arange(half, dtype=torch.float64) * 2.0 / head_dim)
+    return t.to(device)
+def rope_neox(x, pos, base=None, inv=None):
+    """NEOX rope over pairs (i, i + d/2), fp32; x [T, H, D], pos [T]; pass
+    `inv` to reuse a precomputed table."""
+    T, H, D = x.shape
+    half = D // 2
+    if inv is None:
+        inv = rope_inv_freq(D, base, x.device)
+    ang = pos.to(torch.float64)[:, None] * inv[None, :]          # [T, half]
+    cos = ang.cos().to(torch.float32)[:, None, :]
+    sin = ang.sin().to(torch.float32)[:, None, :]
+    xf = x.float()
+    x1, x2 = xf[..., :half], xf[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+class DFlashRef:
+    """Reference drafter over an fp32 weight dict (layout as in
+    `random_weights`) with a sliding ring KV state."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cpu",
+                 ring_dtype="fp32"):
+        self.cfg = cfg
+        self.w = {k: v.to(device=device, dtype=torch.float32)
+                  for k, v in weights.items()}
+        self.device = torch.device(device)
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base, device)
+        # per-layer ring K/V [window, Hkv, D]; ring_pos = slot's position, -1 empty
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int64,
+                                   device=device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        """Fuse per-position target features [T, n_feat] to g [T, n_embd]."""
+        g = feats.float() @ self.w["fc"]
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Project fused features into every layer's ring at `pos` [T], chunked
+        at `window` so no two tokens in one pass share a slot."""
+        C = self.cfg
+        for beg in range(0, len(pos), C.window):
+            p = pos[beg:beg + C.window]
+            gc = g[beg:beg + C.window]
+            slots = p % C.window
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+                k = rope_neox(k, p, inv=self.rope_inv)
+                self.ring_k[i, slots] = k.to(self.ring_dt)
+                self.ring_v[i, slots] = v.to(self.ring_dt)
+            self.ring_pos[slots] = p
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        """Denoise [id_last, mask x n_draft] at n_past..; -> [n_draft+1, n_vocab]."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = id_last
+        pos = n_past + torch.arange(B, device=self.device)
+        x = self.w["tok_embd"][tokens]
+        live = self.ring_pos >= 0
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head, C.head_dim)
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            q = rms_norm(q, self.w[f"q_norm.{i}"], C.rms_eps)
+            k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+            q = rope_neox(q, pos, inv=self.rope_inv)
+            k = rope_neox(k, pos, inv=self.rope_inv)
+            # keys = live ring + block; non-causal, SWA masks only the deep past
+            keys = torch.cat([self.ring_k[i, live].float(), k], dim=0)
+            vals = torch.cat([self.ring_v[i, live].float(), v], dim=0)
+            kpos = torch.cat([self.ring_pos[live], pos], dim=0)
+            gq = C.n_head // C.n_head_kv
+            qh = q.permute(1, 0, 2)                                # [H, B, D]
+            kh = keys.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            vh = vals.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            att = qh @ kh.transpose(1, 2) / (C.head_dim ** 0.5)    # [H, B, S]
+            masked = (pos[:, None] - kpos[None, :]) >= C.window    # SWA
+            att = att.masked_fill(masked[None, :, :], float("-inf"))
+            att = att.softmax(dim=-1)
+            out = (att @ vh).permute(1, 0, 2).reshape(B, -1)       # [B, H*D]
+            x = x + out @ self.w[f"wo.{i}"]
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + (gate * (f @ self.w[f"ffn_up.{i}"])) @ self.w[f"ffn_down.{i}"]
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return x @ self.w["lm_head"]
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        """Greedy drafts from block positions 1..n_draft (position 0 unused)."""
+        logits = self.block_step(id_last, n_past, n_draft)
+        return logits[1:].argmax(dim=-1)
+def random_weights(cfg: DFlashConfig, seed=0, device="cpu"):
+    """Gaussian weights at transformer-typical scale, for structural tests."""
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    def t(*shape, s=0.02):
+        return torch.randn(*shape, generator=g) * s
+    C = cfg
+    w = {"fc": t(C.n_feat, C.n_embd), "enc_norm": torch.ones(C.n_embd),
+         "out_norm": torch.ones(C.n_embd),
+         "tok_embd": t(C.n_vocab, C.n_embd),
+         "lm_head": t(C.n_embd, C.n_vocab)}
+    for i in range(C.n_layer):
+        w[f"attn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"wq.{i}"] = t(C.n_embd, C.n_head * C.head_dim)
+        w[f"wk.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wv.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wo.{i}"] = t(C.n_head * C.head_dim, C.n_embd)
+        w[f"q_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"k_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"ffn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"ffn_gate.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_up.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_down.{i}"] = t(C.n_ff, C.n_embd)
+    return {k: v.to(device) for k, v in w.items()}

build/torch213-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""DFlash block-diffusion drafter ops; the eager reference defines the
+semantics and the CUDA ops are certified against it."""
+from .reference import DFlashConfig, DFlashRef, random_weights  # noqa: F401
+from .drafter import DFlashDrafter  # noqa: F401
+try:  # built variants ship _ops
+    from ._ops import ops  # noqa: F401
+except ImportError:
+    try:  # a local JIT build registers the same ops under torch.ops.dflash_ext
+        import torch as _torch
+        _torch.ops.dflash_ext.dflash_accept
+        ops = _torch.ops.dflash_ext
+    except Exception:
+        ops = None
+try:
+    from .gguf_io import load_drafter  # noqa: F401
+except ImportError:  # gguf-py absent; the reference stays usable
+    load_drafter = None
+__all__ = ["DFlashConfig", "DFlashRef", "DFlashDrafter", "random_weights",
+           "load_drafter", "ops"]

build/torch213-cxx11-cu132-x86_64-linux/_dflash_ops_cuda_03c9586.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:1277c1710a3490d86bc886565d2074d1a6f6257b01a42b4e7cd3e268f39fc039
+size 1402096

build/torch213-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _dflash_ops_cuda_03c9586
+ops = torch.ops._dflash_ops_cuda_03c9586
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_dflash_ops_cuda_03c9586::{op_name}"

build/torch213-cxx11-cu132-x86_64-linux/drafter.py ADDED Viewed

	@@ -0,0 +1,236 @@

+"""GPU DFlash drafter: cuBLAS GEMMs plus the fused ops, with a CUDA-graphed
+block step."""
+import torch
+from .reference import DFlashConfig, rms_norm, rope_inv_freq
+def _ops():
+    from . import ops
+    if ops is None:
+        raise RuntimeError(
+            "dflash-ops compiled kernels are unavailable; build them with "
+            "load_local.load() in a source checkout, or install a published "
+            "variant with kernels.get_kernel")
+    return ops
+class DFlashDrafter:
+    """Ring-cached DFlash drafter on CUDA. Weight keys match `DFlashRef`."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cuda",
+                 ring_dtype="fp32", dtype="fp32"):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        # dtype sets GEMM weights and activations; residual, norms, softmax stay fp32
+        self.wdt = torch.bfloat16 if dtype == "bf16" else torch.float32
+        self.w = {}
+        for k, v in weights.items():
+            dt = self.wdt if v.dim() == 2 else torch.float32
+            self.w[k] = v.to(device=self.device, dtype=dt).contiguous()
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        # host-built fp64 rope table, shared by every op call
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base,
+                                      self.device).contiguous()
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=self.device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int32,
+                                   device=self.device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        g = (feats.to(self.device, self.wdt) @ self.w["fc"]).float()
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Ring injection at `pos`, chunked at `window` so no slot is written twice per launch."""
+        C = self.cfg
+        o = _ops()
+        pos = pos.to(self.device, torch.int32).contiguous()
+        for beg in range(0, pos.numel(), C.window):
+            p = pos[beg:beg + C.window].contiguous()
+            gc = g[beg:beg + C.window].to(self.wdt)
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv,
+                                                  C.head_dim).contiguous()
+                o.dflash_inject(self.ring_k[i], self.ring_v[i], self.ring_pos,
+                                k, v, self.w[f"k_norm.{i}"], self.rope_inv, p,
+                                C.rms_eps)
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = int(id_last)
+        bpos = (n_past + torch.arange(B, device=self.device)).to(torch.int32)
+        return self._forward(tokens, bpos.contiguous())
+    def _forward(self, tokens, bpos):
+        """Block forward on caller-owned device tensors; reads no host state."""
+        C = self.cfg
+        o = _ops()
+        B = tokens.numel()
+        x = self.w["tok_embd"][tokens].float()
+        scale = C.head_dim ** -0.5
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps).to(self.wdt)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head,
+                                             C.head_dim).contiguous()
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv,
+                                             C.head_dim).contiguous()
+            o.dflash_qk_rope(q, k, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, bpos,
+                             C.rms_eps)
+            att = torch.empty_like(q)
+            o.dflash_block_attn(att, q, k, v, self.ring_k[i], self.ring_v[i],
+                                self.ring_pos, bpos, scale)
+            x = x + (att.reshape(B, -1) @ self.w[f"wo.{i}"]).float()
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps).to(self.wdt)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + ((gate * (f @ self.w[f"ffn_up.{i}"]))
+                     @ self.w[f"ffn_down.{i}"]).float()
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return (x.to(self.wdt) @ self.w["lm_head"]).float()
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        return self.block_step(id_last, n_past, n_draft)[1:].argmax(dim=-1)
+    # ---- graphed step ---------------------------------------------------
+    def _static_buffers(self, B):
+        """Every tensor a captured step touches, allocated once and held."""
+        C = self.cfg
+        d = self.device
+        wdt = self.wdt
+        s = {
+            "tokens": torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                                 device=d),
+            "arange": torch.arange(B, dtype=torch.int32, device=d),
+            "n_past": torch.zeros((), dtype=torch.int32, device=d),
+            "bpos": torch.empty(B, dtype=torch.int32, device=d),
+            "xw": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "x": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "sq": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "ms": torch.empty(B, 1, dtype=torch.float32, device=d),
+            "hn": torch.empty(B, C.n_embd, dtype=torch.float32, device=d),
+            "h": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "q": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "k": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "v": torch.empty(B, C.n_head_kv * C.head_dim, dtype=wdt, device=d),
+            "att": torch.empty(B, C.n_head * C.head_dim, dtype=wdt, device=d),
+            "ao": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "g1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "sg": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "u1": torch.empty(B, C.n_ff, dtype=wdt, device=d),
+            "d1": torch.empty(B, C.n_embd, dtype=wdt, device=d),
+            "lg": torch.empty(B, C.n_vocab, dtype=wdt, device=d),
+            "logits": torch.empty(B, C.n_vocab, dtype=torch.float32,
+                                  device=d),
+        }
+        return s
+    def _rms_static(self, src, w, s, out):
+        """RMS norm of `src` into `out` using only held buffers."""
+        torch.pow(src, 2, out=s["sq"])
+        torch.mean(s["sq"], -1, keepdim=True, out=s["ms"])
+        s["ms"].add_(self.cfg.rms_eps).rsqrt_()
+        torch.mul(src, s["ms"], out=out)
+        out.mul_(w)
+    def _forward_static(self, s):
+        """Allocation-free block forward into held buffers; numerically
+        identical to `_forward`."""
+        C = self.cfg
+        o = _ops()
+        B = s["tokens"].numel()
+        scale = C.head_dim ** -0.5
+        torch.index_select(self.w["tok_embd"], 0, s["tokens"], out=s["xw"])
+        s["x"].copy_(s["xw"])
+        for i in range(C.n_layer):
+            self._rms_static(s["x"], self.w[f"attn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"wq.{i}"], out=s["q"])
+            torch.matmul(s["h"], self.w[f"wk.{i}"], out=s["k"])
+            torch.matmul(s["h"], self.w[f"wv.{i}"], out=s["v"])
+            qv = s["q"].view(B, C.n_head, C.head_dim)
+            kv_ = s["k"].view(B, C.n_head_kv, C.head_dim)
+            vv = s["v"].view(B, C.n_head_kv, C.head_dim)
+            o.dflash_qk_rope(qv, kv_, self.w[f"q_norm.{i}"],
+                             self.w[f"k_norm.{i}"], self.rope_inv, s["bpos"],
+                             C.rms_eps)
+            av = s["att"].view(B, C.n_head, C.head_dim)
+            o.dflash_block_attn(av, qv, kv_, vv, self.ring_k[i],
+                                self.ring_v[i], self.ring_pos, s["bpos"],
+                                scale)
+            torch.matmul(s["att"], self.w[f"wo.{i}"], out=s["ao"])
+            s["x"].add_(s["ao"])
+            self._rms_static(s["x"], self.w[f"ffn_norm.{i}"], s, s["hn"])
+            s["h"].copy_(s["hn"])
+            torch.matmul(s["h"], self.w[f"ffn_gate.{i}"], out=s["g1"])
+            # fused silu keeps the step bitwise-equal to the eager path
+            torch.ops.aten.silu.out(s["g1"], out=s["sg"])
+            torch.matmul(s["h"], self.w[f"ffn_up.{i}"], out=s["u1"])
+            s["sg"].mul_(s["u1"])
+            torch.matmul(s["sg"], self.w[f"ffn_down.{i}"], out=s["d1"])
+            s["x"].add_(s["d1"])
+        self._rms_static(s["x"], self.w["out_norm"], s, s["hn"])
+        s["h"].copy_(s["hn"])
+        torch.matmul(s["h"], self.w["lm_head"], out=s["lg"])
+        s["logits"].copy_(s["lg"])
+        return s["logits"]
+    def capture_step(self, n_draft=None):
+        """Graph-captured block step over static buffers (the region allocates
+        nothing); returns step(id_last, n_past) -> logits, overwritten per replay."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        s = self._static_buffers(n_draft + 1)
+        stream = torch.cuda.Stream(device=self.device)
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):                    # warm kernels and cuBLAS
+                torch.add(s["arange"], s["n_past"], out=s["bpos"])
+                self._forward_static(s)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            torch.add(s["arange"], s["n_past"], out=s["bpos"])  # in-graph: replay sees updates
+            logits = self._forward_static(s)
+        def step(id_last, n_past):
+            s["tokens"][0].fill_(int(id_last))
+            s["n_past"].fill_(int(n_past))
+            graph.replay()
+            return logits
+        step.graph = graph
+        step.buffers = s
+        return step
+    # ---- verification ---------------------------------------------------
+    def accept(self, draft, target):
+        """Leading agreement count, device-side."""
+        o = _ops()
+        n = torch.zeros(1, dtype=torch.int32, device=self.device)
+        o.dflash_accept(n, draft.to(torch.int32).contiguous(),
+                        target.to(torch.int32).contiguous())
+        return n

build/torch213-cxx11-cu132-x86_64-linux/gguf_io.py ADDED Viewed

	@@ -0,0 +1,61 @@

+"""Load the DFlash drafter GGUF into reference-layout fp32 tensors; 2-D weights
+transpose from GGUF [out, in] to [in, out]; tok_embd and lm_head are the
+caller's (they belong to the target model)."""
+import numpy as np
+import torch
+from gguf import GGUFReader
+from gguf import quants as gq
+from .reference import DFlashConfig
+_BLK = {
+    "attn_norm": "attn_norm", "attn_q": "wq", "attn_k": "wk", "attn_v": "wv",
+    "attn_output": "wo", "attn_q_norm": "q_norm", "attn_k_norm": "k_norm",
+    "ffn_norm": "ffn_norm", "ffn_gate": "ffn_gate", "ffn_up": "ffn_up",
+    "ffn_down": "ffn_down",
+}
+_TOP = {"fc": "fc", "enc.output_norm": "enc_norm", "output_norm": "out_norm"}
+def _field(reader, key, default=None):
+    f = reader.get_field(key)
+    if f is None:
+        return default
+    v = f.contents()
+    return v
+def load_drafter(path):
+    """-> (DFlashConfig, weights dict) with fp32 tensors in reference layout."""
+    r = GGUFReader(str(path))
+    cfg = DFlashConfig(
+        n_layer=int(_field(r, "dflash.block_count")),
+        n_embd=int(_field(r, "dflash.embedding_length")),
+        n_head=int(_field(r, "dflash.attention.head_count")),
+        n_head_kv=int(_field(r, "dflash.attention.head_count_kv")),
+        head_dim=int(_field(r, "dflash.attention.key_length")),
+        n_ff=int(_field(r, "dflash.feed_forward_length")),
+        rope_base=float(_field(r, "dflash.rope.freq_base")),
+        rms_eps=float(_field(r, "dflash.attention.layer_norm_rms_epsilon")),
+        window=int(_field(r, "dflash.attention.sliding_window")),
+        block_size=int(_field(r, "dflash.block_size")),
+        target_layers=tuple(_field(r, "dflash.target_layers")),
+    )
+    w = {}
+    for t in r.tensors:
+        name = t.name
+        assert name.endswith(".weight"), name
+        base = name[: -len(".weight")]
+        if base.startswith("blk."):
+            _, idx, part = base.split(".", 2)
+            key = f"{_BLK[part]}.{idx}"
+        else:
+            key = _TOP[base]
+        data = gq.dequantize(t.data, t.tensor_type)
+        arr = torch.from_numpy(np.ascontiguousarray(data)).float()
+        if arr.dim() == 2:
+            arr = arr.t().contiguous()
+        w[key] = arr
+    expect = 3 + 11 * cfg.n_layer
+    assert len(w) == expect, (len(w), expect)
+    return cfg, w

build/torch213-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "dflash-ops",
+  "id": "_dflash_ops_cuda_03c9586",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "wNpuIbQEXTfFNjrmz6aISM46VFgMKAHkqqu+ECJJS20=",
+      "_dflash_ops_cuda_03c9586.abi3.so": "EnfBcQo0kNhryIZWXSB00ab2JXsBpCtOfNPiaPOfwDk=",
+      "_ops.py": "+6e3Oo6lwRkZaeOrhSSZEwWYs4/Fpd/zQgJj4MVEi9I=",
+      "drafter.py": "wGU5TWlXWr/ZogxQr4pho5Wku66UUygnczpsTu82dls=",
+      "gguf_io.py": "XrgF2dH7LNR1KTNKKu/pKhh6wT90bCjcjCAgmyGsuZA=",
+      "reference.py": "InmFS9rUwURZ7ZGVUxwQJCCjJZR/hKWrF3eMfp4zQvY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "03c9586154cc2afa6efaa7b111dccef715ec9d69",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/reference.py ADDED Viewed

	@@ -0,0 +1,188 @@

+"""Eager fp32 DFlash drafter reference, transcribed from llama.cpp's
+llama_model_dflash graph; the CUDA ops are certified against this module."""
+from dataclasses import dataclass, field
+import torch
+@dataclass
+class DFlashConfig:
+    n_layer: int = 5
+    n_embd: int = 6656
+    n_head: int = 32
+    n_head_kv: int = 8
+    head_dim: int = 128
+    n_ff: int = 19968
+    n_vocab: int = 202048
+    rope_base: float = 500000.0
+    rms_eps: float = 1e-6
+    window: int = 2048            # SWA width; also the ring capacity
+    block_size: int = 16
+    mask_token_id: int = 201818
+    target_layers: tuple = (2, 14, 26, 38, 50)
+    @property
+    def n_feat(self):
+        return len(self.target_layers) * self.n_embd
+    @classmethod
+    def tiny(cls):
+        """Toy sizes with the real structure; head_dim stays a multiple of 32."""
+        return cls(n_layer=2, n_embd=64, n_head=4, n_head_kv=2, head_dim=32,
+                   n_ff=96, n_vocab=997, window=32, block_size=8,
+                   mask_token_id=990, target_layers=(1, 2))
+def rms_norm(x, w, eps):
+    """llama.cpp LLM_NORM_RMS: x * w / sqrt(mean(x^2) + eps), fp32."""
+    v = x.float()
+    return (v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps)) * w.float()
+def rope_inv_freq(head_dim, base, device="cpu"):
+    """fp64 NEOX inverse frequencies base^(-2i/head_dim), host-evaluated so the
+    table is device-invariant."""
+    half = head_dim // 2
+    t = base ** (-torch.arange(half, dtype=torch.float64) * 2.0 / head_dim)
+    return t.to(device)
+def rope_neox(x, pos, base=None, inv=None):
+    """NEOX rope over pairs (i, i + d/2), fp32; x [T, H, D], pos [T]; pass
+    `inv` to reuse a precomputed table."""
+    T, H, D = x.shape
+    half = D // 2
+    if inv is None:
+        inv = rope_inv_freq(D, base, x.device)
+    ang = pos.to(torch.float64)[:, None] * inv[None, :]          # [T, half]
+    cos = ang.cos().to(torch.float32)[:, None, :]
+    sin = ang.sin().to(torch.float32)[:, None, :]
+    xf = x.float()
+    x1, x2 = xf[..., :half], xf[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+class DFlashRef:
+    """Reference drafter over an fp32 weight dict (layout as in
+    `random_weights`) with a sliding ring KV state."""
+    def __init__(self, cfg: DFlashConfig, weights: dict, device="cpu",
+                 ring_dtype="fp32"):
+        self.cfg = cfg
+        self.w = {k: v.to(device=device, dtype=torch.float32)
+                  for k, v in weights.items()}
+        self.device = torch.device(device)
+        C = cfg
+        # ring storage dtype; arithmetic stays fp32
+        self.ring_dt = (torch.bfloat16 if ring_dtype == "bf16"
+                        else torch.float32)
+        self.rope_inv = rope_inv_freq(C.head_dim, C.rope_base, device)
+        # per-layer ring K/V [window, Hkv, D]; ring_pos = slot's position, -1 empty
+        self.ring_k = torch.zeros(C.n_layer, C.window, C.n_head_kv, C.head_dim,
+                                  device=device, dtype=self.ring_dt)
+        self.ring_v = torch.zeros_like(self.ring_k)
+        self.ring_pos = torch.full((C.window,), -1, dtype=torch.int64,
+                                   device=device)
+    # ---- encoder --------------------------------------------------------
+    def encode(self, feats):
+        """Fuse per-position target features [T, n_feat] to g [T, n_embd]."""
+        g = feats.float() @ self.w["fc"]
+        return rms_norm(g, self.w["enc_norm"], self.cfg.rms_eps)
+    # ---- KV injection ---------------------------------------------------
+    def inject(self, g, pos):
+        """Project fused features into every layer's ring at `pos` [T], chunked
+        at `window` so no two tokens in one pass share a slot."""
+        C = self.cfg
+        for beg in range(0, len(pos), C.window):
+            p = pos[beg:beg + C.window]
+            gc = g[beg:beg + C.window]
+            slots = p % C.window
+            for i in range(C.n_layer):
+                k = (gc @ self.w[f"wk.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                v = (gc @ self.w[f"wv.{i}"]).view(-1, C.n_head_kv, C.head_dim)
+                k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+                k = rope_neox(k, p, inv=self.rope_inv)
+                self.ring_k[i, slots] = k.to(self.ring_dt)
+                self.ring_v[i, slots] = v.to(self.ring_dt)
+            self.ring_pos[slots] = p
+    # ---- noise-block forward -------------------------------------------
+    def block_step(self, id_last, n_past, n_draft=None):
+        """Denoise [id_last, mask x n_draft] at n_past..; -> [n_draft+1, n_vocab]."""
+        C = self.cfg
+        if n_draft is None:
+            n_draft = C.block_size - 1
+        B = n_draft + 1
+        tokens = torch.full((B,), C.mask_token_id, dtype=torch.int64,
+                            device=self.device)
+        tokens[0] = id_last
+        pos = n_past + torch.arange(B, device=self.device)
+        x = self.w["tok_embd"][tokens]
+        live = self.ring_pos >= 0
+        for i in range(C.n_layer):
+            h = rms_norm(x, self.w[f"attn_norm.{i}"], C.rms_eps)
+            q = (h @ self.w[f"wq.{i}"]).view(B, C.n_head, C.head_dim)
+            k = (h @ self.w[f"wk.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            v = (h @ self.w[f"wv.{i}"]).view(B, C.n_head_kv, C.head_dim)
+            q = rms_norm(q, self.w[f"q_norm.{i}"], C.rms_eps)
+            k = rms_norm(k, self.w[f"k_norm.{i}"], C.rms_eps)
+            q = rope_neox(q, pos, inv=self.rope_inv)
+            k = rope_neox(k, pos, inv=self.rope_inv)
+            # keys = live ring + block; non-causal, SWA masks only the deep past
+            keys = torch.cat([self.ring_k[i, live].float(), k], dim=0)
+            vals = torch.cat([self.ring_v[i, live].float(), v], dim=0)
+            kpos = torch.cat([self.ring_pos[live], pos], dim=0)
+            gq = C.n_head // C.n_head_kv
+            qh = q.permute(1, 0, 2)                                # [H, B, D]
+            kh = keys.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            vh = vals.permute(1, 0, 2).repeat_interleave(gq, dim=0)
+            att = qh @ kh.transpose(1, 2) / (C.head_dim ** 0.5)    # [H, B, S]
+            masked = (pos[:, None] - kpos[None, :]) >= C.window    # SWA
+            att = att.masked_fill(masked[None, :, :], float("-inf"))
+            att = att.softmax(dim=-1)
+            out = (att @ vh).permute(1, 0, 2).reshape(B, -1)       # [B, H*D]
+            x = x + out @ self.w[f"wo.{i}"]
+            f = rms_norm(x, self.w[f"ffn_norm.{i}"], C.rms_eps)
+            gate = torch.nn.functional.silu(f @ self.w[f"ffn_gate.{i}"])
+            x = x + (gate * (f @ self.w[f"ffn_up.{i}"])) @ self.w[f"ffn_down.{i}"]
+        x = rms_norm(x, self.w["out_norm"], C.rms_eps)
+        return x @ self.w["lm_head"]
+    def draft_greedy(self, id_last, n_past, n_draft=None):
+        """Greedy drafts from block positions 1..n_draft (position 0 unused)."""
+        logits = self.block_step(id_last, n_past, n_draft)
+        return logits[1:].argmax(dim=-1)
+def random_weights(cfg: DFlashConfig, seed=0, device="cpu"):
+    """Gaussian weights at transformer-typical scale, for structural tests."""
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    def t(*shape, s=0.02):
+        return torch.randn(*shape, generator=g) * s
+    C = cfg
+    w = {"fc": t(C.n_feat, C.n_embd), "enc_norm": torch.ones(C.n_embd),
+         "out_norm": torch.ones(C.n_embd),
+         "tok_embd": t(C.n_vocab, C.n_embd),
+         "lm_head": t(C.n_embd, C.n_vocab)}
+    for i in range(C.n_layer):
+        w[f"attn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"wq.{i}"] = t(C.n_embd, C.n_head * C.head_dim)
+        w[f"wk.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wv.{i}"] = t(C.n_embd, C.n_head_kv * C.head_dim)
+        w[f"wo.{i}"] = t(C.n_head * C.head_dim, C.n_embd)
+        w[f"q_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"k_norm.{i}"] = torch.ones(C.head_dim)
+        w[f"ffn_norm.{i}"] = torch.ones(C.n_embd)
+        w[f"ffn_gate.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_up.{i}"] = t(C.n_embd, C.n_ff)
+        w[f"ffn_down.{i}"] = t(C.n_ff, C.n_embd)
+    return {k: v.to(device) for k, v in w.items()}