Upload folder using huggingface_hub

configuration_luau_coder.py

@@ -1,146 +1,154 @@
-from __future__ import annotations
-
-import math
-
-from transformers import PretrainedConfig
-
-__all__ = ["LuauCoderConfig"]
-
-
-class LuauCoderConfig(PretrainedConfig):
-    model_type = "luau_coder"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size: int = 32768,
-        hidden_size: int = 2560,
-        intermediate_size: int = 6144,
-        num_hidden_layers: int = 36,
-        num_attention_heads: int = 20,
-        num_key_value_heads: int = 5,
-        head_dim: int = 128,
-        rms_norm_eps: float = 1e-6,
-        final_logit_softcapping: float | None = None,
-        max_position_embeddings: int = 4096,
-        rope_theta: float = 500000.0,
-        rope_scaling: dict | None = None,
-        lerope_alpha: list[float] | None = None,
-        nope_layers: list[int] | None = None,
-        diff_v2: bool = True,
-        xsa: bool = True,
-        qk_norm: bool = True,
-        use_output_gate: bool | str | None = False,
-        attn_res: bool = False,
-        attn_res_block_size: int = 6,
-        attn_res_delta_block: bool = False,
-        attn_res_partial_source: bool = False,
-        ccgqa: bool = False,
-        ccgqa_seq_kernel: int = 2,
-        ccgqa_head_kernel: int = 2,
-        ccgqa_qk_mean: bool = True,
-        ccgqa_value_shift: bool = True,
-        vha: bool = False,
-        iha: bool = False,
-        iha_pseudo_heads: int = 2,
-        iha_window: int = 0,
-        iha_global_every: int = 5,
-        initializer_range: float | None = None,
-        use_cache: bool = True,
-        tie_word_embeddings: bool = False,
-        # throughput knobs (no effect on the math)
-        fused_kernels: bool = True,
-        compile_elementwise: bool = True,
-        loss_chunk_size: int = 4096,
-        bos_token_id: int | None = None,
-        eos_token_id: int | list[int] | None = 32744,
-        pad_token_id: int | None = 32745,
-        **kwargs,
-    ):
-        if kwargs.pop("msar", False):
-            raise NotImplementedError("Block MSAR checkpoints are not supported by this modeling code")
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim
-        self.rms_norm_eps = rms_norm_eps
-        self.final_logit_softcapping = final_logit_softcapping  # logits -> c * tanh(logits / c)
-        self.max_position_embeddings = max_position_embeddings
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        # LeRoPE (arXiv:2607.10134): learned log-frequency scale per RoPE band (head_dim / 2), shared by
-        # every layer; frequency_m = exp(alpha_m) / theta^(2m / head_dim). None = plain RoPE.
-        self.lerope_alpha = None if lerope_alpha is None else [float(a) for a in lerope_alpha]
-        self.nope_layers = sorted(nope_layers or [])  # attention without RoPE at these layers
-        self.diff_v2 = diff_v2
-        self.xsa = xsa
-        self.qk_norm = qk_norm
-        # Gated Attention: False | "elementwise" (W_g's rows in qkv) | "headwise" (gate_proj)
-        self.use_output_gate = ("elementwise" if use_output_gate is True
-                                else str(use_output_gate).lower() if use_output_gate else False)
-        self.attn_res = attn_res
-        self.attn_res_block_size = attn_res_block_size
-        self.attn_res_delta_block = attn_res_delta_block
-        self.attn_res_partial_source = attn_res_partial_source
-        self.ccgqa = bool(ccgqa)
-        self.ccgqa_seq_kernel = ccgqa_seq_kernel
-        self.ccgqa_head_kernel = ccgqa_head_kernel
-        self.ccgqa_qk_mean = ccgqa_qk_mean
-        self.ccgqa_value_shift = ccgqa_value_shift
-        self.vha = bool(vha)
-        self.iha = bool(iha)
-        self.iha_pseudo_heads = iha_pseudo_heads
-        self.iha_window = iha_window
-        self.iha_global_every = iha_global_every
-        if self.iha and (iha_pseudo_heads < 1 or iha_window < 0 or iha_global_every < 0):
-            raise ValueError("iha: iha_pseudo_heads >= 1, iha_window >= 0 and iha_global_every >= 0")
-        # training default ("auto"): 1/sqrt(hidden_size)
-        self.initializer_range = hidden_size ** -0.5 if initializer_range is None else initializer_range
-        self.use_cache = use_cache
-        self.fused_kernels = fused_kernels
-        self.compile_elementwise = compile_elementwise
-        self.loss_chunk_size = loss_chunk_size
-        if num_attention_heads % num_key_value_heads:
-            raise ValueError(f"num_attention_heads ({num_attention_heads}) must be a multiple of "
-                             f"num_key_value_heads ({num_key_value_heads})")
-        if head_dim % 2:
-            raise ValueError("head_dim must be even (RoPE)")
-        if self.use_output_gate not in (False, "elementwise", "headwise"):
-            raise ValueError(f"use_output_gate: true / 'elementwise', 'headwise' or false, got {use_output_gate!r}")
-        if self.ccgqa and not qk_norm:
-            raise ValueError("ccgqa needs qk_norm")
-        if self.ccgqa and ccgqa_value_shift and (num_key_value_heads * head_dim) % 2:
-            raise ValueError("ccgqa_value_shift needs an even num_key_value_heads * head_dim")
-        if attn_res and math.ceil(num_hidden_layers / attn_res_block_size) > 32:
-            raise ValueError("attn_res: at most 32 blocks (num_hidden_layers / attn_res_block_size)")
-        if any(not 0 <= i < num_hidden_layers for i in self.nope_layers):
-            raise ValueError(f"nope_layers must be in [0, {num_hidden_layers}), got {self.nope_layers}")
-        if rope_scaling is not None:
-            rtype = rope_scaling.get("rope_type", rope_scaling.get("type"))
-            if rtype not in (None, "default", "yarn"):
-                raise ValueError(f"rope_scaling: only 'yarn' is supported, got {rtype!r}")
-        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id,
-                         pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    @property
-    def output_gate(self) -> str | None:
-        """None (no gate), "elementwise" or "headwise"."""
-        return self.use_output_gate or None
-
-    @property
-    def num_query_heads(self) -> int:
-        """Query heads the attention kernel runs (2x the output heads with DiffAttn V2)."""
-        return self.num_attention_heads * (2 if self.diff_v2 else 1)
-
-    @property
-    def num_q_latent_heads(self) -> int:
-        """Query heads ``qkv`` produces (VHA: the physical heads every KV group shares)."""
-        return self.num_query_heads // self.num_key_value_heads if self.vha else self.num_query_heads
-
-    def iha_layer(self, i: int) -> bool:
-        """Whether layer ``i`` runs IHA (the rest are plain global attention layers)."""
-        every = self.iha_global_every
-        return self.iha and not (every and (i + 1) % every == 0)
+"""LuauCoder configuration (``trust_remote_code=True``).
+
+Mirrors ``lm.config.ModelConfig`` with transformers' standard names.
+``num_attention_heads`` counts attention output heads: with ``diff_v2`` the model
+runs twice as many query heads against ``num_key_value_heads`` K/V heads.
+"""
+
+from __future__ import annotations
+
+import math
+
+from transformers import PretrainedConfig
+
+__all__ = ["LuauCoderConfig"]
+
+
+class LuauCoderConfig(PretrainedConfig):
+    model_type = "luau_coder"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 32768,
+        hidden_size: int = 2560,
+        intermediate_size: int = 6144,
+        num_hidden_layers: int = 36,
+        num_attention_heads: int = 20,
+        num_key_value_heads: int = 5,
+        head_dim: int = 128,
+        rms_norm_eps: float = 1e-6,
+        final_logit_softcapping: float | None = None,
+        max_position_embeddings: int = 4096,
+        rope_theta: float = 500000.0,
+        rope_scaling: dict | None = None,
+        lerope_alpha: list[float] | None = None,
+        nope_layers: list[int] | None = None,
+        diff_v2: bool = True,
+        xsa: bool = True,
+        qk_norm: bool = True,
+        use_output_gate: bool | str | None = False,
+        attn_res: bool = False,
+        attn_res_block_size: int = 6,
+        attn_res_delta_block: bool = False,
+        attn_res_partial_source: bool = False,
+        ccgqa: bool = False,
+        ccgqa_seq_kernel: int = 2,
+        ccgqa_head_kernel: int = 2,
+        ccgqa_qk_mean: bool = True,
+        ccgqa_value_shift: bool = True,
+        vha: bool = False,
+        iha: bool = False,
+        iha_pseudo_heads: int = 2,
+        iha_window: int = 0,
+        iha_global_every: int = 5,
+        head_gate_rank: int = 0,
+        initializer_range: float | None = None,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = False,
+        # throughput knobs (no effect on the math)
+        fused_kernels: bool = True,
+        compile_elementwise: bool = True,
+        loss_chunk_size: int = 4096,
+        bos_token_id: int | None = None,
+        eos_token_id: int | list[int] | None = 32744,
+        pad_token_id: int | None = 32745,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.rms_norm_eps = rms_norm_eps
+        self.final_logit_softcapping = final_logit_softcapping
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        # LeRoPE: learned log-frequency scale per RoPE band, shared by every layer; None = plain RoPE
+        self.lerope_alpha = None if lerope_alpha is None else [float(a) for a in lerope_alpha]
+        self.nope_layers = sorted(nope_layers or [])
+        self.diff_v2 = diff_v2
+        self.xsa = xsa
+        self.qk_norm = qk_norm
+        # False | "elementwise" | "headwise"
+        self.use_output_gate = ("elementwise" if use_output_gate is True
+                                else str(use_output_gate).lower() if use_output_gate else False)
+        self.attn_res = attn_res
+        self.attn_res_block_size = attn_res_block_size
+        self.attn_res_delta_block = attn_res_delta_block
+        self.attn_res_partial_source = attn_res_partial_source
+        self.ccgqa = bool(ccgqa)
+        self.ccgqa_seq_kernel = ccgqa_seq_kernel
+        self.ccgqa_head_kernel = ccgqa_head_kernel
+        self.ccgqa_qk_mean = ccgqa_qk_mean
+        self.ccgqa_value_shift = ccgqa_value_shift
+        self.vha = bool(vha)
+        self.iha = bool(iha)
+        self.iha_pseudo_heads = iha_pseudo_heads
+        self.iha_window = iha_window
+        self.iha_global_every = iha_global_every
+        # gated byte head: [h, sigmoid(u . h + b) * P h] against lm_head = [W, Z]; 0 = plain head
+        self.head_gate_rank = int(head_gate_rank)
+        if self.head_gate_rank < 0:
+            raise ValueError("head_gate_rank must be >= 0")
+        if self.iha and (iha_pseudo_heads < 1 or iha_window < 0 or iha_global_every < 0):
+            raise ValueError("iha: iha_pseudo_heads >= 1, iha_window >= 0 and iha_global_every >= 0")
+        self.initializer_range = hidden_size ** -0.5 if initializer_range is None else initializer_range
+        self.use_cache = use_cache
+        self.fused_kernels = fused_kernels
+        self.compile_elementwise = compile_elementwise
+        self.loss_chunk_size = loss_chunk_size
+        if num_attention_heads % num_key_value_heads:
+            raise ValueError(f"num_attention_heads ({num_attention_heads}) must be a multiple of "
+                             f"num_key_value_heads ({num_key_value_heads})")
+        if head_dim % 2:
+            raise ValueError("head_dim must be even (RoPE)")
+        if self.use_output_gate not in (False, "elementwise", "headwise"):
+            raise ValueError(f"use_output_gate: true / 'elementwise', 'headwise' or false, got {use_output_gate!r}")
+        if self.ccgqa and not qk_norm:
+            raise ValueError("ccgqa needs qk_norm")
+        if self.ccgqa and ccgqa_value_shift and (num_key_value_heads * head_dim) % 2:
+            raise ValueError("ccgqa_value_shift needs an even num_key_value_heads * head_dim")
+        if attn_res and math.ceil(num_hidden_layers / attn_res_block_size) > 32:
+            raise ValueError("attn_res: at most 32 blocks (num_hidden_layers / attn_res_block_size)")
+        if any(not 0 <= i < num_hidden_layers for i in self.nope_layers):
+            raise ValueError(f"nope_layers must be in [0, {num_hidden_layers}), got {self.nope_layers}")
+        if rope_scaling is not None:
+            rtype = rope_scaling.get("rope_type", rope_scaling.get("type"))
+            if rtype not in (None, "default", "yarn"):
+                raise ValueError(f"rope_scaling: only 'yarn' is supported, got {rtype!r}")
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id,
+                         pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+    @property
+    def output_gate(self) -> str | None:
+        """None (no gate), "elementwise" or "headwise"."""
+        return self.use_output_gate or None
+
+    @property
+    def num_query_heads(self) -> int:
+        """Query heads the attention kernel runs (2x the output heads with DiffAttn V2)."""
+        return self.num_attention_heads * (2 if self.diff_v2 else 1)
+
+    @property
+    def num_q_latent_heads(self) -> int:
+        """Query heads ``qkv`` produces (VHA: the physical heads every KV group shares)."""
+        return self.num_query_heads // self.num_key_value_heads if self.vha else self.num_query_heads
+
+    def iha_layer(self, i: int) -> bool:
+        """Whether layer ``i`` runs IHA (the rest are plain global attention layers)."""
+        every = self.iha_global_every
+        return self.iha and not (every and (i + 1) % every == 0)

luau_gemv_tuned.py

@@ -1,46 +1,119 @@
-TABLE = {'NVIDIA RTX PRO 6000 Blackwell Server Edition': {'arch': [12, 0, 188],
-                                                  'gemv': {'1,1408,1280,0,0': [4, 512, 4, 4, 0],
-                                                           '1,1408,1280,1,0': [4, 1024, 2, 4, 0],
-                                                           '1,1408,3840,0,0': [2, 1024, 2, 2, 0],
-                                                           '1,1408,3840,1,0': [2, 1024, 2, 4, 0],
-                                                           '1,1802,1408,0,0': [2, 512, 1, 1, 0],
-                                                           '1,1802,1408,1,0': [2, 512, 1, 3, 0],
-                                                           '1,32768,1408,0,0': [4, 256, 2, 4, 0],
-                                                           '1,3840,1408,0,1': [4, 256, 1, 3, 0],
-                                                           '1,3840,1408,1,1': [4, 512, 1, 3, 0],
-                                                           '16,1408,1280,0,0': [32, 256, 2, 4, 1],
-                                                           '16,1408,1280,1,0': [32, 256, 4, 4, 1],
-                                                           '16,1408,3840,0,0': [16, 256, 2, 4, 1],
-                                                           '16,1408,3840,1,0': [32, 256, 4, 4, 1],
-                                                           '16,1802,1408,0,0': [16, 128, 2, 4, 1],
-                                                           '16,1802,1408,1,0': [32, 128, 4, 4, 1],
-                                                           '16,32768,1408,0,0': [128, 64, 4, 4, 1],
-                                                           '16,3840,1408,0,1': [32, 256, 4, 4, 1],
-                                                           '16,3840,1408,1,1': [32, 128, 4, 4, 1],
-                                                           '2,1408,1280,0,0': [4, 512, 2, 4, 0],
-                                                           '2,1408,1280,1,0': [4, 1024, 2, 2, 0],
-                                                           '2,1408,3840,0,0': [2, 1024, 2, 1, 0],
-                                                           '2,1408,3840,1,0': [4, 1024, 2, 2, 0],
-                                                           '2,1802,1408,0,0': [1, 512, 2, 1, 0],
-                                                           '2,1802,1408,1,0': [2, 512, 2, 2, 0],
-                                                           '2,32768,1408,0,0': [128, 64, 4, 4, 1],
-                                                           '2,3840,1408,0,1': [16, 256, 2, 3, 1],
-                                                           '2,3840,1408,1,1': [2, 512, 1, 3, 0],
-                                                           '4,1408,1280,0,0': [4, 256, 2, 2, 0],
-                                                           '4,1408,1280,1,0': [4, 256, 2, 3, 0],
-                                                           '4,1408,3840,0,0': [2, 1024, 2, 2, 0],
-                                                           '4,1408,3840,1,0': [4, 1024, 4, 1, 0],
-                                                           '4,1802,1408,0,0': [2, 512, 2, 3, 0],
-                                                           '4,1802,1408,1,0': [2, 512, 2, 4, 0],
-                                                           '4,32768,1408,0,0': [64, 64, 4, 4, 1],
-                                                           '4,3840,1408,0,1': [16, 256, 2, 3, 1],
-                                                           '4,3840,1408,1,1': [32, 256, 4, 3, 1],
-                                                           '8,1408,1280,0,0': [16, 256, 2, 4, 1],
-                                                           '8,1408,1280,1,0': [32, 256, 4, 4, 1],
-                                                           '8,1408,3840,0,0': [16, 256, 2, 4, 1],
-                                                           '8,1408,3840,1,0': [4, 512, 4, 4, 0],
-                                                           '8,1802,1408,0,0': [32, 256, 4, 3, 1],
-                                                           '8,1802,1408,1,0': [32, 128, 4, 4, 1],
-                                                           '8,32768,1408,0,0': [128, 64, 4, 4, 1],
-                                                           '8,3840,1408,0,1': [32, 256, 4, 3, 1],
-                                                           '8,3840,1408,1,1': [32, 128, 4, 4, 1]}}}
+"""Pinned decode-GEMV configs per GPU (written by ``benchmark/tune_gemv.py``).
+
+``TABLE[gpu name] = {"arch": [major, minor, SMs], "gemv": {key: (BLOCK_N, BLOCK_K, warps, stages, DOT)}}``;
+a GPU missing here uses an entry of the same architecture and SM count (e.g. the
+Workstation and Server editions of a card), else the kernels' defaults.
+"""
+
+TABLE = {'NVIDIA GeForce RTX 4060': {'arch': [8, 9, 24],
+                             'gemv': {'1,1024,1024,0,0': [1, 1024, 1, 3, 0],
+                                      '1,1024,1024,1,0': [16, 256, 4, 4, 0],
+                                      '1,1024,3072,0,0': [1, 1024, 4, 2, 0],
+                                      '1,1024,3072,1,0': [2, 512, 1, 1, 0],
+                                      '1,1408,1280,0,0': [32, 128, 1, 1, 0],
+                                      '1,1408,1280,1,0': [8, 256, 1, 2, 0],
+                                      '1,1408,3840,0,0': [1, 1024, 1, 2, 0],
+                                      '1,1408,3840,1,0': [4, 1024, 1, 1, 0],
+                                      '1,1422,1792,0,0': [1, 1024, 1, 1, 0],
+                                      '1,1422,1792,1,0': [8, 256, 1, 1, 0],
+                                      '1,1544,1024,0,0': [1, 512, 1, 1, 0],
+                                      '1,1544,1024,1,0': [4, 1024, 2, 2, 0],
+                                      '1,1792,1792,0,0': [4, 256, 1, 4, 0],
+                                      '1,1792,1792,1,0': [4, 256, 1, 2, 0],
+                                      '1,1792,4864,0,0': [4, 256, 1, 1, 0],
+                                      '1,1792,4864,1,0': [2, 1024, 1, 1, 0],
+                                      '1,1802,1408,0,0': [8, 64, 1, 2, 0],
+                                      '1,1802,1408,1,0': [16, 128, 1, 2, 0],
+                                      '1,2048,2048,0,0': [2, 1024, 2, 2, 0],
+                                      '1,2048,2048,1,0': [32, 128, 2, 3, 0],
+                                      '1,2048,6144,0,0': [2, 512, 1, 1, 0],
+                                      '1,2048,6144,1,0': [4, 512, 1, 3, 0],
+                                      '1,257,1792,0,0': [1, 1024, 1, 1, 0],
+                                      '1,257,1792,1,0': [1, 1024, 1, 4, 0],
+                                      '1,2592,2048,0,0': [2, 512, 1, 2, 0],
+                                      '1,2592,2048,1,0': [2, 512, 4, 4, 0],
+                                      '1,3072,1024,0,1': [1, 1024, 1, 3, 0],
+                                      '1,3072,1024,1,1': [4, 1024, 1, 1, 0],
+                                      '1,32768,1024,0,0': [2, 512, 1, 1, 0],
+                                      '1,32768,1408,0,0': [1, 1024, 2, 1, 0],
+                                      '1,32768,2048,0,0': [4, 512, 4, 1, 0],
+                                      '1,3840,1408,0,1': [2, 1024, 1, 4, 0],
+                                      '1,3840,1408,1,1': [8, 512, 1, 3, 0],
+                                      '1,4864,1792,0,1': [1, 1024, 1, 1, 0],
+                                      '1,4864,1792,1,1': [16, 256, 1, 1, 0],
+                                      '1,6144,2048,0,1': [1, 1024, 1, 3, 0],
+                                      '1,6144,2048,1,1': [2, 1024, 1, 3, 0]}},
+ 'NVIDIA RTX PRO 6000 Blackwell Server Edition': {'arch': [12, 0, 188],
+                                                  'gemv': {'1,1024,1024,0,0': [2, 1024, 1, 2, 0],
+                                                           '1,1024,1024,1,0': [8, 1024, 4, 1, 0],
+                                                           '1,1024,3072,0,0': [2, 1024, 2, 2, 0],
+                                                           '1,1024,3072,1,0': [2, 1024, 4, 1, 0],
+                                                           '1,1408,1280,0,0': [4, 512, 4, 3, 0],
+                                                           '1,1408,1280,1,0': [4, 1024, 2, 3, 0],
+                                                           '1,1408,3840,0,0': [2, 1024, 2, 2, 0],
+                                                           '1,1408,3840,1,0': [2, 1024, 2, 2, 0],
+                                                           '1,1422,1792,0,0': [1, 1024, 4, 1, 0],
+                                                           '1,1422,1792,1,0': [2, 1024, 1, 3, 0],
+                                                           '1,1544,1024,0,0': [1, 1024, 2, 2, 0],
+                                                           '1,1544,1024,1,0': [2, 1024, 1, 4, 0],
+                                                           '1,1792,1792,0,0': [2, 1024, 2, 4, 0],
+                                                           '1,1792,1792,1,0': [2, 1024, 1, 1, 0],
+                                                           '1,1792,4864,0,0': [2, 1024, 2, 3, 0],
+                                                           '1,1792,4864,1,0': [2, 1024, 1, 2, 0],
+                                                           '1,1802,1408,0,0': [2, 512, 1, 3, 0],
+                                                           '1,1802,1408,1,0': [2, 512, 1, 2, 0],
+                                                           '1,2048,2048,0,0': [2, 1024, 2, 2, 0],
+                                                           '1,2048,2048,1,0': [2, 1024, 1, 2, 0],
+                                                           '1,2048,6144,0,0': [2, 1024, 4, 3, 0],
+                                                           '1,2048,6144,1,0': [2, 1024, 1, 3, 0],
+                                                           '1,257,1792,0,0': [4, 1024, 4, 3, 0],
+                                                           '1,257,1792,1,0': [4, 1024, 4, 4, 0],
+                                                           '1,2592,2048,0,0': [2, 512, 1, 1, 0],
+                                                           '1,2592,2048,1,0': [2, 1024, 1, 2, 0],
+                                                           '1,3072,1024,0,1': [4, 256, 1, 3, 0],
+                                                           '1,3072,1024,1,1': [4, 1024, 2, 1, 0],
+                                                           '1,32768,1024,0,0': [2, 256, 1, 1, 0],
+                                                           '1,32768,1408,0,0': [1, 1024, 4, 2, 0],
+                                                           '1,32768,2048,0,0': [8, 256, 1, 2, 0],
+                                                           '1,3840,1408,0,1': [2, 512, 1, 2, 0],
+                                                           '1,3840,1408,1,1': [4, 512, 1, 3, 0],
+                                                           '1,4864,1792,0,1': [4, 256, 1, 4, 0],
+                                                           '1,4864,1792,1,1': [4, 1024, 1, 2, 0],
+                                                           '1,6144,2048,0,1': [2, 256, 1, 3, 0],
+                                                           '1,6144,2048,1,1': [2, 512, 1, 3, 0],
+                                                           '16,1408,1280,0,0': [32, 256, 2, 4, 1],
+                                                           '16,1408,1280,1,0': [32, 256, 4, 4, 1],
+                                                           '16,1408,3840,0,0': [16, 256, 2, 4, 1],
+                                                           '16,1408,3840,1,0': [32, 256, 4, 4, 1],
+                                                           '16,1802,1408,0,0': [16, 128, 2, 4, 1],
+                                                           '16,1802,1408,1,0': [32, 128, 4, 4, 1],
+                                                           '16,32768,1408,0,0': [128, 64, 4, 4, 1],
+                                                           '16,3840,1408,0,1': [32, 256, 4, 4, 1],
+                                                           '16,3840,1408,1,1': [32, 128, 4, 4, 1],
+                                                           '2,1408,1280,0,0': [4, 512, 2, 4, 0],
+                                                           '2,1408,1280,1,0': [4, 1024, 2, 2, 0],
+                                                           '2,1408,3840,0,0': [2, 1024, 2, 1, 0],
+                                                           '2,1408,3840,1,0': [4, 1024, 2, 2, 0],
+                                                           '2,1802,1408,0,0': [1, 512, 2, 1, 0],
+                                                           '2,1802,1408,1,0': [2, 512, 2, 2, 0],
+                                                           '2,32768,1408,0,0': [128, 64, 4, 4, 1],
+                                                           '2,3840,1408,0,1': [16, 256, 2, 3, 1],
+                                                           '2,3840,1408,1,1': [2, 512, 1, 3, 0],
+                                                           '4,1408,1280,0,0': [4, 256, 2, 2, 0],
+                                                           '4,1408,1280,1,0': [4, 256, 2, 3, 0],
+                                                           '4,1408,3840,0,0': [2, 1024, 2, 2, 0],
+                                                           '4,1408,3840,1,0': [4, 1024, 4, 1, 0],
+                                                           '4,1802,1408,0,0': [2, 512, 2, 3, 0],
+                                                           '4,1802,1408,1,0': [2, 512, 2, 4, 0],
+                                                           '4,32768,1408,0,0': [64, 64, 4, 4, 1],
+                                                           '4,3840,1408,0,1': [16, 256, 2, 3, 1],
+                                                           '4,3840,1408,1,1': [32, 256, 4, 3, 1],
+                                                           '8,1408,1280,0,0': [16, 256, 2, 4, 1],
+                                                           '8,1408,1280,1,0': [32, 256, 4, 4, 1],
+                                                           '8,1408,3840,0,0': [16, 256, 2, 4, 1],
+                                                           '8,1408,3840,1,0': [4, 512, 4, 4, 0],
+                                                           '8,1802,1408,0,0': [32, 256, 4, 3, 1],
+                                                           '8,1802,1408,1,0': [32, 128, 4, 4, 1],
+                                                           '8,32768,1408,0,0': [128, 64, 4, 4, 1],
+                                                           '8,3840,1408,0,1': [32, 256, 4, 3, 1],
+                                                           '8,3840,1408,1,1': [32, 128, 4, 4, 1]}}}

luau_infer.py

@@ -1,6 +1,13 @@
+"""Fused decoding engine for the LuauCoder HF port.
+
+~9 Triton kernels per layer (``luau_kernels.py``), one CUDA graph per step, and
+prompt-lookup speculative decoding that keeps the model's own output.
+"""
+
 from __future__ import annotations
 
 import math
+import os
 import sys
 import time
 from collections import deque
@@ -11,7 +18,7 @@ from torch import Tensor
 
 from .luau_kernels import kernels as K
 
-K_ = K   # (``K`` is also the draft length in the decode code)
+K_ = K
 
 E4M3 = torch.float8_e4m3fn
 
@@ -19,7 +26,10 @@ E4M3 = torch.float8_e4m3fn
 def supported(config) -> str | None:
     """None if the engine can run this config, else why not."""
     if getattr(config, "iha", False):
-        return "IHA layers (no decode kernels yet)"
+        if config.diff_v2:
+            return "IHA with DiffAttn V2"
+        if not 1 <= config.iha_pseudo_heads <= 4:
+            return "IHA with more than 4 pseudo-heads"
     if not config.qk_norm:
         return "needs qk_norm"
     if config.num_query_heads // config.num_key_value_heads > 64:
@@ -46,12 +56,17 @@ class _Linear:
         if quant == "fp8":
             self.w8, self.scale = _q8(self.w)
 
+    @property
+    def weight(self) -> Tensor:
+        """What the GEMV reads (the L2 prefetch target)."""
+        return self.w8 if self.w8 is not None else self.w
+
     def gemv(self, x: Tensor, y: Tensor, epi: int = 0, n: int | None = None, up_off: int = 0,
-             cap: float = 0.0) -> None:
+             cap: float = 0.0, pf: Tensor | None = None) -> None:
         if self.w8 is not None:
-            K.gemv(x, self.w8, y, epi, self.scale, n, up_off, cap)
+            K.gemv(x, self.w8, y, epi, self.scale, n, up_off, cap, pf=pf)
         else:
-            K.gemv(x, self.w, y, epi, None, n, up_off, cap)
+            K.gemv(x, self.w, y, epi, None, n, up_off, cap, pf=pf)
 
 
 class _Layer:
@@ -59,16 +74,17 @@ class _Layer:
 
 
 class Engine:
-    """Caches, work buffers and CUDA graphs for decoding one ``LuauCoderForCausalLM``: ``prefill``
-    a (left-padded) prompt batch, then ``step`` a decode graph (``graph(K, settings)``)."""
+    """Caches, work buffers and CUDA graphs for decoding one ``LuauCoderForCausalLM``."""
 
-    MAX_GEMV_ROWS = 16     # rows the skinny GEMV takes (tensor-core tiles from 4); more go to cuBLAS
-    NSPL = 64              # flash-decoding spans per (row, KV head); their length follows the context
+    MAX_GEMV_ROWS = 16  # rows the skinny GEMV takes; more go to cuBLAS
+    NSPL = 32
     L_ROUND = 256          # cache length granularity
-    DECODE_T = 16          # tokens per row the decode kernels take; longer calls (prompts) take the prompt path
-    MAX_DRAFT = DECODE_T - 1   # a verify is K+1 tokens per row (also < the latent ring - 2)
+    DECODE_T = 16  # tokens per row of the decode kernels; longer calls take the prompt path
+    MAX_DRAFT = DECODE_T - 1
 
-    @torch.no_grad()   # (quantising / concatenating parameters must not build autograd graphs)
+    MAX_TILE_ROWS = 128
+
+    @torch.no_grad()
     def __init__(self, model, quant: str | None = None):
         cfg = self.config = model.config
         why = supported(cfg)
@@ -84,12 +100,12 @@ class Engine:
         self.F = cfg.intermediate_size
         self.eps = cfg.rms_norm_eps
         self.diff, self.xsa = cfg.diff_v2, cfg.xsa
-        # output gate: 0 off, 1 per channel (W_g rows in qkv), 2 per head (gate_proj rows appended)
+        # output gate: 0 off, 1 per channel, 2 per head
         self.gate = {"elementwise": 1, "headwise": 2}.get(cfg.output_gate, 0)
         self.gate_w = {1: self.nh * self.hd, 2: self.nh}.get(self.gate, 0)   # gate logits per row
         self.ccgqa = cfg.ccgqa
         self.attn_res = cfg.attn_res
-        self.delta = bool(cfg.attn_res and cfg.attn_res_delta_block)   # Delta Block: mix + stream
+        self.delta = bool(cfg.attn_res and cfg.attn_res_delta_block)
         self.partial = not self.delta or bool(cfg.attn_res_partial_source)   # the partial is routed
         self.S_blk = cfg.attn_res_block_size if self.attn_res else 0
         self.n_src = (math.ceil(self.L / self.S_blk) + 1) if self.attn_res else 1
@@ -99,14 +115,31 @@ class Engine:
         self.gate_off = self.v_off + self.nkv * self.hd
         self.lam_off = self.gate_off + self.gate_w
         self.n_qkv = self.lam_off + (self.nh if self.diff else 0)
+        self.iha_p = cfg.iha_pseudo_heads if getattr(cfg, "iha", False) else 0
+        self.win = cfg.iha_window * self.iha_p if self.iha_p else 0
+        P1 = max(self.iha_p, 1)
+        self.decode_t = (min(self.DECODE_T, self.MAX_TILE_ROWS // (P1 * (self.nq // self.nkv)))
+                         if self.iha_p else self.DECODE_T)
+        self.max_draft = self.decode_t - 1
+        self.NSPL = int(os.environ.get("LUAU_NSPL", self.NSPL))
+        # IHA layers see at most win pseudo-keys: few long spans
+        self.win_nspl = int(os.environ.get("LUAU_WIN_NSPL", self.NSPL if not self.win else 16))
+        self.win_bl = int(os.environ.get("LUAU_WIN_BL", 128))
+        self.attn_bl = int(os.environ.get("LUAU_ATTN_BL", 64))
+        self.win_warps = int(os.environ.get("LUAU_WIN_WARPS", 0)) or None
         self.embed = m.embed.weight
         self.head = _Linear(model.lm_head.weight, None)
+        hg = getattr(model, "head_gate", None)
+        self.hg_rank = hg.rank if hg is not None else 0
+        if hg is not None:
+            self.hg_proj = _Linear(hg.proj.to(self.dtype), None)
+            self.hg_bias = hg.bias.detach().float().contiguous()
         self.final_g = m.norm.weight.float().contiguous()
         f32 = lambda t: t.detach().float().contiguous()
         self.layers: list[_Layer] = []
         for i, blk in enumerate(m.blocks):
             a, ly = blk.attn, _Layer()
-            w = a.qkv.weight   # q | k | v (| per-channel W_g) (| per-head W_g) (| W_lambda)
+            w = a.qkv.weight  # q | k | v (| W_g) (| W_lambda)
             if a.gate_proj is not None:
                 w = torch.cat((w, a.gate_proj.weight), 0)
             if self.diff:
@@ -115,14 +148,22 @@ class Engine:
             ly.w13, ly.w2 = _Linear(blk.ffn.w13.weight, quant), _Linear(blk.ffn.w2.weight, quant)
             ly.attn_g, ly.ffn_g = f32(blk.attn_norm.weight), f32(blk.ffn_norm.weight)
             ly.rope = a.rope
+            ly.iha = bool(getattr(a, "iha", False))
             ly.qg = f32(a.q_norm.weight)
             kg = a.k_norm.weight.float()[None].expand(self.nkv, self.hd)
-            if self.ccgqa:
+            if ly.iha:
+                bf = lambda t: t.detach().to(self.dtype).contiguous()
+                ly.aq, ly.ak, ly.av, ly.ar = bf(a.iha_q), bf(a.iha_k), bf(a.iha_v), bf(a.iha_r)
+                ly.kg2 = f32(a.k_norm.weight)
+                ly.temp = (a.k_temp.float().exp() if self.ccgqa
+                           else torch.ones(self.nkv, device=self.dev)).contiguous()
+            elif self.ccgqa:
                 kg = kg * a.k_temp.float().exp()[:, None]
+            if self.ccgqa:
                 ly.conv = tuple(t.detach().contiguous() for t in
                                 (a.conv_seq, a.conv_seq_bias, a.conv_head, a.conv_head_bias))
                 ly.ks, ly.kc = a.conv_seq.shape[1], a.conv_head.shape[1]
-            else:
+            if not self.ccgqa:
                 ly.conv, ly.ks, ly.kc = None, 1, 1
             ly.kg = kg.contiguous()
             if self.attn_res:
@@ -142,7 +183,6 @@ class Engine:
         self.B = self.Lc = 0
         self.graphs: dict = {}
 
-    # buffers
 
     def _rope_tables(self, n: int) -> None:
         if self.cos is None or self.cos.shape[0] < n:
@@ -151,27 +191,33 @@ class Engine:
             self.cos, self.sin = cos[0].contiguous(), sin[0].contiguous()
 
     def allocate(self, B: int, L: int) -> None:
-        """Caches for ``B`` sequences of up to ``L`` tokens (kept while big enough:
-        CUDA graphs hold the addresses)."""
+        """Caches for ``B`` sequences of up to ``L`` tokens (kept while big enough for the CUDA graphs)."""
         L = -(-L // self.L_ROUND) * self.L_ROUND
         if B == self.B and L <= self.Lc:
             return
         L = max(L, self.Lc if B == self.B else 0)
-        self.graphs.clear()   # free the old graphs' pools and caches before the new ones (8 GB cards)
-        self.kc = self.vc = self.lat = self.vh = self.hist = self.pres = None
+        self.graphs.clear()  # free the old graphs first (8 GB cards)
+        self.kc = self.vc = self.lat = self.vh = self.hist = self.pres = self.kl = self.vl = None
         torch.cuda.empty_cache()
         self.B, self.Lc = B, L
         dev, dt = self.dev, self.dtype
-        self.kc = [torch.zeros(B, self.nkv, L, self.hd, device=dev, dtype=dt) for _ in range(self.L)]
-        self.vc = [torch.zeros(B, self.nkv, L, self.hd, device=dev, dtype=dt) for _ in range(self.L)]
+        # IHA layers: interleaved pseudo-keys / -values
+        P1 = max(self.iha_p, 1)
+        self.kc = [torch.zeros(B, self.nkv, L * (P1 if ly.iha else 1), self.hd, device=dev, dtype=dt)
+                   for ly in self.layers]
+        self.vc = [torch.zeros(B, self.nkv, L * (P1 if ly.iha else 1), self.hd, device=dev, dtype=dt)
+                   for ly in self.layers]
+        if self.iha_p:
+            self.kl = torch.zeros(B, self.nkv, L, self.hd, device=dev, dtype=dt)
+            self.vl = torch.zeros(B, self.nkv, L, self.hd, device=dev, dtype=dt)
         if self.ccgqa:
-            ring = K.RING.value   # conv latents / value-shift halves: only the last 2 slots are read
+            ring = K.RING.value
             self.lat = [torch.zeros(B, ring, self.nxq + self.nkv, self.hd, device=dev, dtype=dt)
                         for _ in range(self.L)]
             self.vh = [torch.zeros(B, ring, self.nkv * self.hd // 2, device=dev, dtype=dt) for _ in range(self.L)]
         else:
             self.lat = self.vh = [None] * self.L
-        self._rope_tables(L)
+        self._rope_tables(L * P1)
         self.start = torch.zeros(B, dtype=torch.int32, device=dev)
         self.pres = torch.zeros(B, self.config.vocab_size, dtype=torch.bool, device=dev)
         self.done = torch.zeros(B, dtype=torch.bool, device=dev)
@@ -179,65 +225,108 @@ class Engine:
         self.cur = torch.zeros(B, dtype=torch.int32, device=dev)
         self.pcur = torch.zeros(B, dtype=torch.int32, device=dev)
         self.seed = torch.zeros(1, dtype=torch.long, device=dev)
+        self.eos_ids = torch.full((self.MAX_EOS,), -1, dtype=torch.long, device=dev)
+
+    MAX_EOS = 8
+
+    def set_eos(self, eos: tuple[int, ...]) -> None:
+        if len(eos) > self.MAX_EOS:
+            raise ValueError(f"at most {self.MAX_EOS} EOS ids, got {len(eos)}")
+        self.eos_ids.fill_(-1)
+        if eos:
+            self.eos_ids[:len(eos)] = torch.tensor(eos, dtype=torch.long, device=self.dev)
+
+    def _is_eos(self, t: Tensor) -> Tensor:
+        return (t[..., None] == self.eos_ids).any(-1)
 
     def _bufs(self, R: int, decode: bool = True) -> dict:
-        """Work buffers for ``R`` rows (``decode``: with the decode-attention scratch; a prompt
-        over ``DECODE_T`` tokens per row attends through sdpa and needs none)."""
+        """Work buffers for ``R`` rows."""
         dev, dt = self.dev, self.dtype
         G = self.nq // self.nkv
         g_pad = max(16, 1 << (G - 1).bit_length())
         nspl = self.NSPL if decode else 0
-        return dict(
+        P1 = max(self.iha_p, 1)
+        bf = dict(
             S=torch.zeros(R, self.n_src, self.D, device=dev, dtype=dt),
             xn=torch.empty(R, self.D, device=dev, dtype=dt),
             qkv=torch.empty(R, self.n_qkv, device=dev, dtype=dt),
-            q=torch.empty(R, self.nq, self.hd, device=dev, dtype=dt),
-            y=torch.zeros(R, self.nh * self.hd, device=dev, dtype=dt),   # (prefill pad rows: never written)
+            q=torch.empty(R * P1, self.nq, self.hd, device=dev, dtype=dt),
+            y=torch.zeros(R, self.nh * self.hd, device=dev, dtype=dt),
             f=torch.empty(R, self.F, device=dev, dtype=dt),
-            ml=torch.empty(R, self.nkv, nspl, 2, g_pad, device=dev, dtype=torch.float32),
-            acc=torch.empty(R, self.nkv, nspl, g_pad, self.hd, device=dev, dtype=torch.float32),
+            ml=torch.empty(R * P1, self.nkv, nspl, 2, g_pad, device=dev, dtype=torch.float32),
+            acc=torch.empty(R * P1, self.nkv, nspl, g_pad, self.hd, device=dev, dtype=torch.float32),
             cnt=torch.zeros(R, self.nkv, device=dev, dtype=torch.int32),
         )
+        if self.iha_p:
+            bf["ql"] = torch.empty(R, self.nq, self.hd, device=dev, dtype=dt)
+            bf["vslot"] = torch.empty(R * P1, device=dev, dtype=torch.int32)
+            bf["vstart"] = torch.empty(R * P1, device=dev, dtype=torch.int32)
+        if self.hg_rank:
+            bf["haug"] = torch.zeros(R, self.D + self.hg_rank, device=dev, dtype=dt)
+            bf["hgy"] = torch.empty(R, self.hg_rank + 1, device=dev, dtype=dt)
+        return bf
 
-    # forward
 
     def _forward(self, bf: dict, ids: Tensor, pos: Tensor, slot: Tensor, start: Tensor, T: int,
                  last: Tensor | None, logits: Tensor) -> None:
-        """Rows ``R = B * T`` (``T`` consecutive tokens per sequence). Writes the caches
-        and the logits of rows ``last`` (None: all) into ``logits`` (fp32)."""
+        """Forward ``R = B * T`` rows: writes the caches and the fp32 logits of rows ``last`` (None: all)."""
         R = ids.shape[0]
         small = R <= self.MAX_GEMV_ROWS
         S, xn, qkv, q, y, f = bf["S"], bf["xn"], bf["qkv"], bf["q"], bf["y"], bf["f"]
+        P = self.iha_p
+        if P:
+            ar = torch.arange(P, device=self.dev, dtype=torch.int32)
+            torch.add(slot[:, None] * P, ar, out=bf["vslot"][:R * P].view(R, P))
+            bf["vstart"][:R * P].view(R, P).copy_((start * P)[:, None].expand(R, P))
+            vslot, vstart = bf["vslot"][:R * P], bf["vstart"][:R * P]
+        # L2 prefetch: each GEMV's weights stream in while the kernels before it run
+        sites = os.environ.get("LUAU_PF_SITES", "qk,w13").split(",")
+        pfs = {k: ((lambda lin: lin.weight) if small and k in sites else (lambda lin: None))
+               for k in ("mix", "qk", "attn", "w13")}
         for i, ly in enumerate(self.layers):
             # ---- attention input
             if i == 0:
-                K.mix_norm(S, 1, 1, xn, ly.attn_g, self.eps, ids=ids, emb=self.embed)
+                K.mix_norm(S, 1, 1, xn, ly.attn_g, self.eps, ids=ids, emb=self.embed, pf=pfs["mix"](ly.qkv))
             elif self.attn_res:
                 K.mix_norm(S, ly.attn_ne, 2, xn, ly.attn_g, self.eps, w=ly.attn_w, add=self.delta,
-                           partial=self.partial)
+                           partial=self.partial, pf=pfs["mix"](ly.qkv))
             else:
-                K.mix_norm(S, 1, 0, xn, ly.attn_g, self.eps)
+                K.mix_norm(S, 1, 0, xn, ly.attn_g, self.eps, pf=pfs["mix"](ly.qkv))
             self._lin(ly.qkv, xn, qkv, small)
-            if self.ccgqa and T > self.DECODE_T:   # a prompt: token tiles (one conv-weight read per tile)
-                K.qk_prefill(qkv, self.lat[i], self.vh[i], pos, self.cos, self.sin, ly.conv, ly.qg, ly.kg, q,
-                             self.kc[i], self.vc[i], T, self.eps, nq=self.nq, nkv=self.nkv, nxq=self.nxq,
+            ql, kc, vc = (bf["ql"], self.kl, self.vl) if ly.iha else (q, self.kc[i], self.vc[i])
+            rope = ly.rope and not ly.iha
+            if self.ccgqa and T > self.decode_t:
+                K.qk_prefill(qkv, self.lat[i], self.vh[i], pos, self.cos, self.sin, ly.conv, ly.qg, ly.kg, ql,
+                             kc, vc, T, self.eps, nq=self.nq, nkv=self.nkv, nxq=self.nxq,
                              hd=self.hd, v_off=self.v_off, ks=ly.ks, kc_taps=ly.kc,
                              qk_mean=bool(self.config.ccgqa_qk_mean),
-                             vshift=bool(self.config.ccgqa_value_shift), rope=ly.rope)
+                             vshift=bool(self.config.ccgqa_value_shift), rope=rope)
             else:
-                K.qk_decode(qkv, self.lat[i], self.vh[i], pos, slot, self.cos, self.sin, ly.conv, ly.qg, ly.kg, q,
-                            self.kc[i], self.vc[i], T, self.eps, nq=self.nq, nkv=self.nkv, nxq=self.nxq,
+                K.qk_decode(qkv, self.lat[i], self.vh[i], pos, slot, self.cos, self.sin, ly.conv, ly.qg, ly.kg, ql,
+                            kc, vc, T, self.eps, nq=self.nq, nkv=self.nkv, nxq=self.nxq,
                             hd=self.hd, v_off=self.v_off, ks=ly.ks, kc_taps=ly.kc, ccgqa=self.ccgqa,
                             qk_mean=bool(self.ccgqa and self.config.ccgqa_qk_mean),
-                            vshift=bool(self.ccgqa and self.config.ccgqa_value_shift), rope=ly.rope)
-            if T <= self.DECODE_T:
-                K.attention(q, self.kc[i], self.vc[i], slot, start, qkv, y, bf["ml"], bf["acc"], bf["cnt"], T,
+                            vshift=bool(self.ccgqa and self.config.ccgqa_value_shift), rope=rope,
+                            pf=pfs["qk"](ly.o))
+            if ly.iha:
+                K.iha_mix(ql, self.kl, self.vl, ly.aq, ly.ak, ly.av, ly.qg, ly.kg2, ly.temp, pos, slot,
+                          self.cos, self.sin, q, self.kc[i], self.vc[i], T, self.eps, P=P, nq=self.nq,
+                          nkv=self.nkv, hd=self.hd, rope=ly.rope)
+                if T <= self.decode_t:
+                    K.attention(q[:R * P], self.kc[i], self.vc[i], vslot, vstart, qkv, y, bf["ml"], bf["acc"],
+                                bf["cnt"], T, nq=self.nq, nkv=self.nkv, hd=self.hd, nspl=self.win_nspl,
+                                bl=self.win_bl, warps=self.win_warps,
+                                lam_off=self.lam_off, gate_off=self.gate_off, diff=False, xsa=self.xsa,
+                                gate=self.gate, win=self.win, iha_p=P, iha_r=ly.ar, pf=pfs["attn"](ly.w13))
+                else:
+                    self._iha_attention_sdpa(i, ly, q, qkv, y, start, T)
+            elif T <= self.decode_t:
+                K.attention(q[:R], self.kc[i], self.vc[i], slot, start, qkv, y, bf["ml"], bf["acc"], bf["cnt"], T,
                             nq=self.nq, nkv=self.nkv, hd=self.hd, nspl=self.NSPL,
                             lam_off=self.lam_off, gate_off=self.gate_off, diff=self.diff, xsa=self.xsa,
-                            gate=self.gate)
+                            gate=self.gate, pf=pfs["attn"](ly.w13), bl=self.attn_bl)
             else:
-                self._attention_sdpa(i, q, qkv, y, slot, start, T)
-            # ---- attention output -> partial (AttnRes: source slot block+1)
+                self._attention_sdpa(i, q[:R], qkv, y, slot, start, T)
             dst = S[:, ly.block + 1] if self.attn_res else S[:, 0]
             fresh = self.attn_res and ly.block_start
             self._lin(ly.o, y, dst, small, epi=0 if fresh else 1)
@@ -247,36 +336,42 @@ class Engine:
                            partial=self.partial)
             else:
                 K.mix_norm(S, 1, 0, xn, ly.ffn_g, self.eps)
-            self._lin(ly.w13, xn, f, small, epi=2)
+            self._lin(ly.w13, xn, f, small, epi=2, pf=pfs["w13"](ly.w2))
             self._lin(ly.w2, f, dst, small, epi=1)
         # ---- output
         Sl = S if last is None else S[last]
-        xo = xn if last is None else xn[:Sl.shape[0]]
+        n_out = Sl.shape[0]
+        xo = (bf["haug"][:n_out, :self.D] if self.hg_rank else
+              xn if last is None else xn[:n_out])
         if self.delta:
             K.mix_norm(Sl, self.n_src, 3, xo, self.final_g, self.eps)
         elif self.attn_res:
             K.mix_norm(Sl, self.n_src, 2, xo, self.final_g, self.eps, w=self.out_w)
         else:
             K.mix_norm(Sl, 1, 0, xo, self.final_g, self.eps)
+        if self.hg_rank:
+            hy = bf["hgy"][:n_out]
+            self._lin(self.hg_proj, xo, hy, n_out <= self.MAX_GEMV_ROWS)
+            K.head_gate(hy, self.hg_bias, bf["haug"], self.D, self.hg_rank)
+            xo = bf["haug"][:n_out]
         if xo.shape[0] <= self.MAX_GEMV_ROWS:
             self.head.gemv(xo, logits, epi=3, cap=self.cap)
         else:
             z = (xo @ self.head.w.T).float()
             logits.copy_(self.cap * torch.tanh(z / self.cap) if self.cap else z)
 
-    def _lin(self, lin: _Linear, x: Tensor, y: Tensor, small: bool, epi: int = 0) -> None:
-        """``y = x @ W.T`` with the GEMV epilogue ``epi`` (0 store, 1 add, 2 SwiGLU): the skinny
-        GEMV for decode-sized ``small`` calls, else cuBLAS + the epilogue in torch."""
+    def _lin(self, lin: _Linear, x: Tensor, y: Tensor, small: bool, epi: int = 0, pf: Tensor | None = None) -> None:
+        """``y = x @ W.T`` with epilogue ``epi``: skinny GEMV for small calls, else cuBLAS."""
         if small:
             if epi == 2:
-                lin.gemv(x, y, 2, n=self.F, up_off=self.F)
+                lin.gemv(x, y, 2, n=self.F, up_off=self.F, pf=pf)
             else:
-                lin.gemv(x, y, epi)
+                lin.gemv(x, y, epi, pf=pf)
             return
         z = x @ lin.w.T
         if epi == 2:
             g, u = z.chunk(2, -1)
-            torch.mul(F.silu(g), u, out=y)   # bf16 like the port's eager SwiGLU (no fp32 copies)
+            torch.mul(F.silu(g), u, out=y)
         elif epi == 1:
             y.add_(z)
         else:
@@ -289,12 +384,12 @@ class Engine:
         R = q.shape[0]
         B = R // T
         sl = slot.view(B, T).long()
-        n = T   # a prompt: slots 0 .. T-1
+        n = T
         qq = q.view(B, T, self.nq, self.hd).transpose(1, 2)
         kk, vv = self.kc[i][:, :, :n], self.vc[i][:, :, :n]
-        if not self._padded:   # causal flash attention (the masked path can't skip the upper half)
+        if not self._padded:
             o = F.scaled_dot_product_attention(qq, kk, vv, is_causal=True, enable_gqa=True)
-        else:   # left padding: a row's keys from its first real token; pad rows see themselves (no empty rows)
+        else:
             keys = torch.arange(n, device=self.dev)
             mask = ((keys >= start.view(B, T, 1).long()) & (keys <= sl[..., None])) | (keys == sl[..., None])
             o = F.scaled_dot_product_attention(qq, kk, vv, attn_mask=mask[:, None], enable_gqa=True)
@@ -304,12 +399,43 @@ class Engine:
         out = _diff_xsa_fn(o.transpose(1, 2), lam, own, self.nh, self.nkv, self.xsa, gate)
         y.copy_(out.reshape(R, -1))
 
-    # sampling
+    IHA_CHUNK = 2048
+
+    def _iha_attention_sdpa(self, i: int, ly, q: Tensor, qkv: Tensor, y: Tensor, start: Tensor, T: int) -> None:
+        """Prefill attention of an IHA layer, in bands of ``IHA_CHUNK`` pseudo-rows."""
+        port = sys.modules[type(self.model).__module__]
+        P, nq, nkv, hd = self.iha_p, self.nq, self.nkv, self.hd
+        R = y.shape[0]
+        B = R // T
+        Tv = T * P
+        qq = q[:R * P].view(B, Tv, nq, hd).transpose(1, 2)
+        kk, vv = self.kc[i][:, :, :Tv], self.vc[i][:, :, :Tv]
+        out = torch.empty(B, Tv, nq, hd, device=self.dev, dtype=self.dtype)
+        first = start.view(B, T)[:, 0].long() * P
+        W = self.win
+        for a in range(0, Tv, self.IHA_CHUNK):
+            e = min(Tv, a + self.IHA_CHUNK)
+            k0 = max(0, a - W + 1) if W else 0
+            dq = torch.arange(a, e, device=self.dev)[:, None] - torch.arange(k0, e, device=self.dev)[None]
+            mask = dq >= 0
+            if W:
+                mask = mask & (dq < W)
+            if self._padded:
+                keys = torch.arange(k0, e, device=self.dev)
+                mask = (mask[None] & (keys >= first[:, None, None])) | (dq == 0)[None]
+                mask = mask[:, None]
+            o = F.scaled_dot_product_attention(qq[:, :, a:e], kk[:, :, k0:e], vv[:, :, k0:e], attn_mask=mask,
+                                               enable_gqa=True)
+            out[:, a:e] = o.transpose(1, 2)
+        yc = torch.einsum("btphd,hp->bthd", out.view(B, T, P, nq, hd), ly.ar)
+        vp = vv.reshape(B, nkv, T, P, hd).permute(0, 2, 3, 1, 4)
+        gate = qkv[:, self.gate_off:self.gate_off + self.gate_w].view(B, T, -1) if self.gate else None
+        y.copy_(port._iha_diff_xsa(yc, None, vp, self.nh, nkv, self.xsa, gate).reshape(R, -1))
+
 
     @staticmethod
     def process(logits: Tensor, sp: dict, pres: Tensor | None) -> Tensor:
-        """transformers' logits processors in ``generate()``'s order, in torch: the reference
-        the sampler kernel (``luau_kernels.sample``) is tested against."""
+        """transformers' logits processors in torch (reference for the sampler kernel)."""
         x = logits
         if sp["repetition_penalty"] != 1.0:
             p = sp["repetition_penalty"]
@@ -343,11 +469,7 @@ class Engine:
                   ids=y, T=1)
         return y
 
-    # prefill / decode
-    #
-    # Decode state (device): ``hist`` (B, Lc) every token by cache slot, ``cur`` the slot of
-    # the newest token (sampled, not yet run through the model), ``pcur`` its document
-    # position, ``done`` (EOS seen / stopped by the host: the row is frozen).
+    # decode state: hist (token by cache slot), cur (newest slot), pcur (its position), done
 
     def prefill(self, input_ids: Tensor, attention_mask: Tensor, sp: dict) -> Tensor:
         """Fill the caches with left-padded prompts ``(B, T0)``; returns the first sampled tokens."""
@@ -357,20 +479,20 @@ class Engine:
         self.start.copy_(T0 - am.sum(1))
         self.done.zero_()
         self.pres.zero_()
-        self.seed.fill_(int(torch.randint(1 << 30, (1,))))   # the sampler kernel's Philox seed (manual_seed)
-        if sp["repetition_penalty"] != 1.0:   # every prompt id, pads too (as transformers' processor)
+        self.seed.fill_(int(torch.randint(1 << 30, (1,))))
+        if sp["repetition_penalty"] != 1.0:
             self.pres.scatter_(1, input_ids, True)
         R = B * T0
-        bf = self._bufs(R, decode=T0 <= self.DECODE_T)
+        bf = self._bufs(R, decode=T0 <= self.decode_t)
         pos = (am.cumsum(1) - 1).clamp_min(0).to(torch.int32).reshape(R)
         slot = torch.arange(T0, device=self.dev, dtype=torch.int32).repeat(B)
         start = self.start.repeat_interleave(T0)
         last = torch.arange(B, device=self.dev) * T0 + T0 - 1
         logits = torch.empty(B, self.config.vocab_size, device=self.dev, dtype=torch.float32)
         self._forward(bf, input_ids.reshape(R), pos, slot, start, T0, last, logits)
-        self.last_logits = logits   # (the prompts' next-token logits, for inspection / tests)
+        self.last_logits = logits
         first = self.sample(logits, sp).masked_fill(self.done, sp["pad"])
-        self.done |= first == sp["eos"]
+        self.done |= self._is_eos(first)
         if sp["repetition_penalty"] != 1.0:
             self.pres.scatter_(1, first[:, None], True)
         self.hist[:, :T0] = input_ids
@@ -379,34 +501,28 @@ class Engine:
         self.pcur.copy_(pos.view(B, T0)[:, -1] + 1)
         return first
 
-    NGRAM = 3   # longest suffix the draft looks up
+    NGRAM = 3
 
     def _draft(self, K: int, st: dict) -> Tensor:
-        """Prompt-lookup draft ``(B, K)``: the tokens after the latest earlier occurrence of
-        the longest (up to ``NGRAM``) suffix of the history; no occurrence: anything."""
+        """Prompt-lookup draft ``(B, K)`` from the latest earlier occurrence of the history's suffix."""
         Lc, NG = self.Lc, self.NGRAM
         cur = self.cur.long()[:, None]
         hp = st["hp"]
         hp[:, NG:].copy_(self.hist)
         j = st["jidx"]
-        tails = self.hist.gather(1, (cur - st["ng"]).clamp_min(0))          # tails[:, i] = hist[cur - i]
+        tails = self.hist.gather(1, (cur - st["ng"]).clamp_min(0))
         lo = self.start.long()[:, None]
         m = j < cur
         best = cur[:, 0] - 1
         for i in range(NG):
-            m = m & (hp[:, NG - i:NG - i + Lc] == tails[:, i:i + 1])        # hist[j - i] == hist[cur - i]
+            m = m & (hp[:, NG - i:NG - i + Lc] == tails[:, i:i + 1])
             cand = torch.where(m & (j - i >= lo) & (cur - i >= lo), j, -1).amax(1)
             best = torch.where(cand >= 0, cand, best)
         idx = (best[:, None] + 1 + st["ark"]).clamp_max(cur).clamp(0, Lc - 1)
         return self.hist.gather(1, idx)
 
     def _step(self, st: dict, K: int, sp: dict) -> None:
-        """One decode step: the newest token plus a ``K``-token draft through the model
-        (``T = K+1`` rows per sequence), then accept the agreeing draft prefix and sample
-        the token after it. Greedy: exactly greedy decoding. Sampling: speculative
-        sampling with a point-mass draft (accept ``d`` with prob ``p(d)``, else draw from
-        ``p`` without ``d``): exactly the model's distribution. A repetition penalty counts
-        each verify row's earlier draft tokens as seen, as sequential decoding would."""
+        """One decode step: the newest token plus a ``K``-token draft, then accept and sample."""
         B, T, V, Lc = self.B, K + 1, self.config.vocab_size, self.Lc
         ar, ar64 = st["ar"], st["ar64"]
         cur = self.cur
@@ -417,7 +533,7 @@ class Engine:
         pos = self.pcur[:, None] + ar
         self._forward(st["bf"], ids.reshape(-1), pos.reshape(-1), slot.reshape(-1),
                       self.start.repeat_interleave(T), T, None, st["logits"])
-        if sp["sample"] or sp["repetition_penalty"] != 1.0:   # one Triton kernel per row
+        if sp["sample"] or sp["repetition_penalty"] != 1.0:
             self.seed.add_(1)
             R = B * T
             y, pd, u = st["y"], st["pd"], st["u"]
@@ -445,8 +561,8 @@ class Engine:
             out = y
         frozen = self.done.clone()
         out = torch.where(frozen[:, None] | (ar64 > n[:, None]), sp["pad"], out)
-        self.done |= (out == sp["eos"]).any(1)
-        if sp["repetition_penalty"] != 1.0:   # every token this step emitted
+        self.done |= self._is_eos(out).any(1)
+        if sp["repetition_penalty"] != 1.0:
             self.pres.scatter_(1, torch.where(ar64 <= n[:, None], out, out[:, :1]), True)
         self.hist.scatter_(1, (cur.long()[:, None] + 1 + ar64).clamp_max(Lc - 1), out)
         adv = torch.where(frozen, 0, n + 1).to(torch.int32)
@@ -463,13 +579,12 @@ class Engine:
             t.copy_(s)
 
     def graph(self, K: int, sp: dict) -> dict:
-        """The CUDA graph of a decode step with a ``K``-token draft (built once per
-        cache / sampling settings; warm-up runs are undone)."""
+        """The CUDA graph of a decode step with a ``K``-token draft."""
         key = (K, tuple(sorted(sp.items())))
         if key in self.graphs:
             return self.graphs[key]
-        if K > self.MAX_DRAFT:
-            raise ValueError(f"draft length {K}: at most {self.MAX_DRAFT}")
+        if K > self.max_draft:
+            raise ValueError(f"draft length {K}: at most {self.max_draft}")
         B, T, dev = self.B, K + 1, self.dev
         st = dict(bf=self._bufs(B * T), K=K,
                   logits=torch.empty(B * T, self.config.vocab_size, device=dev, dtype=torch.float32),
@@ -491,7 +606,6 @@ class Engine:
         with torch.cuda.graph(g):
             self._step(st, K, sp)
         st["graph"] = g
-        # time it (for the speculation policy), then undo
         g.replay()
         torch.cuda.synchronize()
         e0, e1 = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
@@ -507,22 +621,15 @@ class Engine:
 
     @staticmethod
     def step(st: dict) -> Tensor:
-        """Replay a step graph: ``(B, K+2)`` = the new tokens (``pad`` after the accepted
-        ones) and the number of accepted draft tokens (-1: the row was frozen)."""
+        """Replay a step graph: ``(B, K+2)`` new tokens and the accepted draft count (-1: frozen)."""
         st["graph"].replay()
         return st["res"]
 
 
-# ---------decoding loop
-
-
 class _Lookup:
-    """Host mirror of the engine's prompt-lookup draft (``Engine._draft``) for one row:
-    how many draft tokens speculation would have accepted at each recent position
-    (greedy: exactly; sampling: the same odds), and from that the tokens per step a
-    run of speculative steps would have made there. It decides when to speculate."""
+    """Host mirror of the engine's draft for one row: estimates whether speculation pays."""
 
-    WINDOW = 96   # recent positions the estimate looks at
+    WINDOW = 96
 
     def __init__(self, K: int, ng: int):
         self.K, self.ng = K, ng
@@ -532,8 +639,7 @@ class _Lookup:
         self.acc: dict[int, int] = {}             # position -> draft tokens it would accept
 
     def add(self, t: int) -> None:
-        """The row's next token: resolves the pending drafts it continues or breaks, then drafts
-        from the new position."""
+        """Advance by the row's next token."""
         for d in self.pending:
             if d[3] and d[1] < len(d[0]) and d[0][d[1]] == t:
                 d[1] += 1
@@ -570,26 +676,17 @@ class _Lookup:
 
 
 class FusedDecoder:
-    """Decoding with ``Engine``: the prompt prefilled by the fused kernels, each decode
-    step one CUDA graph of ~9 Triton kernels per layer with sampling on the device. The
-    host reads tokens up to ``DEPTH`` steps behind the GPU (streaming, stop strings, EOS).
-
-    With ``spec`` K > 0 a step can also verify a K-token prompt-lookup draft (same
-    output: greedy stays greedy, sampling keeps the model's distribution). A K-step
-    costs more than a plain one, so it runs only while the recent history says it
-    pays: expected tokens per step (1 + accepted) above the measured cost ratio
-    (sampled batches: off; ``spec < 0`` forces it at every step)."""
+    """Decoding loop around ``Engine``; speculates only while the recent history says it pays."""
 
     DEPTH = 8
-    MARGIN = 1.1       # speculate when the expected tokens per step beat the step cost ratio by this
-    # sampled batches: their verify steps cost ~2x a plain one and the estimate (each row's own
-    # history, read DEPTH steps late) misjudged them in benchmarks: off unless forced (spec < 0)
+    MARGIN = 1.1
+    # sampled batches: speculation off unless forced (spec < 0)
     MARGIN_BATCH_SAMPLED = float("inf")
-    MIN_CACHE = 4096   # cache length allocated at least for B <= 8 (graphs are reused while it fits)
+    MIN_CACHE = 4096
 
     def __init__(self, model, quant: str | None = None, spec: int = 0):
         self.engine = Engine(model, quant)
-        self.spec, self.always = abs(spec), spec < 0   # negative: speculate at every step (testing)
+        self.spec, self.always = min(abs(spec), self.engine.max_draft), spec < 0  # < 0: always
         self.stats = {}
 
     def generate(self, input_ids: Tensor, attention_mask: Tensor, args, tok, streamer=None) -> Tensor:
@@ -603,18 +700,17 @@ class FusedDecoder:
 
     def decode(self, input_ids: Tensor, attention_mask: Tensor | None, *, max_new_tokens: int,
                sample: bool = False, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0,
-               min_p: float = 0.0, repetition_penalty: float = 1.0, num_samples: int = 1, eos: int | None = None,
-               pad: int | None = None, stop=(), tokenizer=None, streamer=None) -> Tensor:
-        """Prompt + completion ids ``(B * num_samples, T0 + n)`` with ``pad`` after a row's end,
-        like transformers' ``generate()``. Left-padded prompts; ``stop`` strings need ``tokenizer``."""
+               min_p: float = 0.0, repetition_penalty: float = 1.0, num_samples: int = 1, eos=None,
+               pad: int | None = None, stop=(), tokenizer=None, streamer=None, stopping_criteria=None) -> Tensor:
+        """Prompt + completion ids ``(B * num_samples, T0 + n)``, ``pad`` after a row's end."""
         eng = self.engine
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
-        eos = -1 if eos is None else int(eos)
-        pad = int(pad) if pad is not None else max(eos, 0)
+        eos_set = tuple(dict.fromkeys(int(e) for e in ((eos,) if isinstance(eos, int) else (eos or ()))))
+        pad = int(pad) if pad is not None else (eos_set[0] if eos_set else 0)
         stop = list(stop or ())
         sp = dict(sample=bool(sample), temperature=float(temperature), top_k=int(top_k or 0), top_p=float(top_p),
-                  min_p=float(min_p or 0.0), repetition_penalty=float(repetition_penalty), eos=eos, pad=pad)
+                  min_p=float(min_p or 0.0), repetition_penalty=float(repetition_penalty), pad=pad)
         if not sp["sample"]:
             sp.update(temperature=1.0, top_k=0, top_p=1.0, min_p=0.0)
         if num_samples > 1:
@@ -627,13 +723,18 @@ class FusedDecoder:
         margin = self.MARGIN_BATCH_SAMPLED if sp["sample"] and B > 1 else self.MARGIN
         K = self.spec if self.always or margin < float("inf") else 0
         eng.allocate(B, max(T0 + n_new + (self.DEPTH + 3) * (K + 1), self.MIN_CACHE if B <= 8 else 0))
+        eng.set_eos(eos_set)
         first = eng.prefill(input_ids, attention_mask, sp)
+        crit_ids = None
+        if stopping_criteria:
+            crit_ids = torch.full((B, T0 + n_new + K + 1), pad, dtype=input_ids.dtype, device=input_ids.device)
+            crit_ids[:, :T0] = input_ids
         if streamer is not None:
             streamer.put(input_ids.cpu())   # skipped as the prompt
         t_setup = time.perf_counter()
         g0 = eng.graph(0, sp) if n_new > 1 else None
         gk = eng.graph(K, sp) if K and n_new > 1 else None
-        setup = time.perf_counter() - t_setup   # first use per batch size / settings: compile + capture
+        setup = time.perf_counter() - t_setup
         ratio = gk["ms"] / g0["ms"] if gk else 0.0
         looks = None
         if gk:
@@ -641,7 +742,7 @@ class FusedDecoder:
             for b, row in enumerate(input_ids.tolist()):
                 for t in row[int(T0 - attention_mask[b].sum()):]:
                     looks[b].add(t)
-        pinned: dict[int, list[Tensor]] = {}   # per step kind: a ring of pinned host buffers
+        pinned: dict[int, list[Tensor]] = {}
         pushed: dict[int, int] = {}
         queue: deque = deque()
 
@@ -678,12 +779,22 @@ class FusedDecoder:
                     gen[b].append(t)
                     if looks is not None:
                         looks[b].add(t)
-                    done[b] = t == eos or len(gen[b]) >= n_new or (bool(stop) and any(
+                    done[b] = t in eos_set or len(gen[b]) >= n_new or (bool(stop) and any(
                         x in tokenizer.decode(gen[b][-tail:], skip_special_tokens=True) for x in stop))
                     if done[b]:
-                        if t != eos:   # the device doesn't know yet: freeze the row
+                        if t not in eos_set:
                             eng.done[b] = True
                         break
+            if crit_ids is not None and not all(done):
+                n = max(len(g) for g in gen)
+                for b, g in enumerate(gen):
+                    if g:
+                        crit_ids[b, T0:T0 + len(g)] = torch.tensor(g, dtype=crit_ids.dtype)
+                hit = torch.as_tensor(stopping_criteria(crit_ids[:, :T0 + n], None)).reshape(-1).tolist()
+                for b in range(B):
+                    if not done[b] and hit[b if len(hit) > 1 else 0]:
+                        done[b] = True
+                        eng.done[b] = True
             if streamer is not None and len(gen[0]) > streamed:
                 streamer.put(torch.tensor(gen[0][streamed:]))
                 streamed = len(gen[0])

luau_kernels.py

@@ -1,734 +1,966 @@
-from __future__ import annotations
-
-import os
-import sys
-from pathlib import Path
-
-import torch
-
-try:   # (remote code: transformers only requires imports outside try blocks)
-    import triton
-    import triton.language as tl
-except ImportError as exc:  # pragma: no cover
-    raise ImportError("the fused LuauCoder decoding kernels need triton") from exc
-
-from .luau_gemv_tuned import TABLE as _TABLE
-
-LOG2E = 1.4426950408889634
-
-#: Programmatic dependent launch (SM90+): each kernel is launched while the previous one
-#: runs and waits for it in ``_pdl`` (griddepcontrol), hiding the launch gap between the
-#: ~260 small kernels of a decode step. ``LUAU_INFER_PDL=0`` turns it off.
-PDL = (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
-       and os.environ.get("LUAU_INFER_PDL", "1") != "0")
-
-
-@triton.jit
-def _pdl(ON: tl.constexpr):
-    """Wait for the previous kernel's writes, then let the next kernel launch."""
-    if ON:
-        tl.extra.cuda.gdc_wait()
-        tl.extra.cuda.gdc_launch_dependents()
-
-
-# AttnRes mix + RMSNorm
-
-
-@triton.jit
-def _mix_norm_kernel(s_ptr, s_row, s_ent, ids_ptr, emb_ptr, w_ptr, g_ptr, out_ptr, o_row, eps, D,
-                     NE: tl.constexpr, NE_PAD: tl.constexpr, MODE: tl.constexpr, BLOCK: tl.constexpr,
-                     ADD: tl.constexpr = 0, PART: tl.constexpr = 1, PDL: tl.constexpr = False):
-    """MODE 0: h = S[r, 0]; 1: h = embed[ids[r]] (also stored to S[r, 0]);
-    2: h = AttnRes mix of S[r, 0..NE-1] with score weight w (``ADD``: Delta Block,
-    plus their sum; ``PART`` off: entry NE-1, the partial, isn't routed); 3: h =
-    their sum (its output). out = rmsnorm(h) * g."""
-    _pdl(PDL)
-    r = tl.program_id(0).to(tl.int64)
-    offs = tl.arange(0, BLOCK)
-    msk = offs < D
-    if MODE == 1:
-        tok = tl.load(ids_ptr + r).to(tl.int64)
-        hb = tl.load(emb_ptr + tok * D + offs, mask=msk, other=0.0)
-        tl.store(s_ptr + r * s_row + offs, hb, mask=msk)
-        h = hb.to(tl.float32)
-    elif MODE == 0:
-        h = tl.load(s_ptr + r * s_row + offs, mask=msk, other=0.0).to(tl.float32)
-    else:   # every entry in one 2-D tile: one reduction per quantity, not one per entry
-        w = tl.load(w_ptr + offs, mask=msk, other=0.0)
-        j = tl.arange(0, NE_PAD)
-        jm = j < NE
-        v = tl.load(s_ptr + r * s_row + j[:, None] * s_ent + offs[None, :], mask=jm[:, None] & msk[None, :],
-                    other=0.0).to(tl.float32)
-        if MODE == 3:
-            h = tl.sum(v, 0)
-        else:
-            sc = tl.sum(v * w[None, :], 1) * tl.rsqrt(tl.sum(v * v, 1) / D + eps)
-            sc = tl.where(j < NE - 1 + PART, sc, float("-inf"))
-            e = tl.exp(sc - tl.max(sc, 0))
-            p = e / tl.sum(e, 0)
-            h = tl.sum((p + ADD)[:, None] * v, 0)
-        h = h.to(out_ptr.dtype.element_ty).to(tl.float32)   # the mix output is bf16
-    g = tl.load(g_ptr + offs, mask=msk, other=0.0).to(tl.float32)
-    y = h * tl.rsqrt(tl.sum(h * h, 0) / D + eps) * g
-    tl.store(out_ptr + r * o_row + offs, y.to(out_ptr.dtype.element_ty), mask=msk)
-
-
-def mix_norm(S: torch.Tensor, n_ent: int, mode: int, out: torch.Tensor, g: torch.Tensor, eps: float,
-             w: torch.Tensor | None = None, ids: torch.Tensor | None = None,
-             emb: torch.Tensor | None = None, add: bool = False, partial: bool = True) -> None:
-    """``S`` ``(R, NS, D)`` sources (row-major per row), ``out`` ``(R, D)``."""
-    R, _, D = S.shape
-    block = triton.next_power_of_2(D)
-    ne_pad = triton.next_power_of_2(n_ent) if mode >= 2 else 1
-    warps = max(4, min(16, ne_pad * block // 1024))
-    _mix_norm_kernel[(R,)](S, S.stride(0), S.stride(1), ids if ids is not None else S,
-                           emb if emb is not None else S, w if w is not None else g, g, out,
-                           out.stride(0), eps, D, NE=n_ent, NE_PAD=ne_pad, MODE=mode, BLOCK=block,
-                           ADD=int(add), PART=int(partial), num_warps=warps, PDL=PDL, launch_pdl=PDL)
-
-
-# skinny GEMM
-
-
-_TUNED_PATH = Path(__file__).with_name("luau_gemv_tuned.py")
-_TUNED: dict | None = None
-
-
-def _tuned() -> dict:
-    """Pinned GEMV configs for this GPU (``benchmark/tune_gemv.py``): its own entry, else one
-    of the same architecture and SM count, else {}."""
-    global _TUNED
-    if _TUNED is None:
-        name = torch.cuda.get_device_name()
-        p = torch.cuda.get_device_properties(0)
-        arch = [p.major, p.minor, p.multi_processor_count]
-        hit = _TABLE.get(name) or next((e for e in _TABLE.values() if e.get("arch") == arch), None)
-        _TUNED = dict(hit["gemv"]) if hit else {}
-    return _TUNED
-
-
-def gemv_key(m_pad: int, n: int, k: int, fp8: bool, epi: int) -> str:
-    return f"{m_pad},{n},{k},{int(fp8)},{int(epi == 2)}"
-
-
-def gemv_config(m_pad: int, n: int, k: int, fp8: bool, epi: int) -> tuple[int, int, int, int, int]:
-    """``(BLOCK_N, BLOCK_K, num_warps, num_stages, DOT)``: pinned, else a bandwidth-friendly
-    default (``DOT``: tensor-core tiles, rows padded to 16; the CUDA-core FMAs of the
-    broadcast path become the limit from ~4 rows)."""
-    hit = _tuned().get(gemv_key(m_pad, n, k, fp8, epi))
-    if hit:
-        return tuple(hit) + (0,) * (5 - len(hit))
-    if m_pad >= 4:
-        return 32, 128, 4, 3, 1
-    bk = 512 if fp8 else 256
-    bn = max(1, min(16, 2048 // (bk * m_pad) * 4))
-    return bn, bk, 4, 2, 0
-
-
-@triton.jit
-def _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn, nm, K, M_PAD: tl.constexpr, BLOCK_N: tl.constexpr,
-               BLOCK_K: tl.constexpr, DOT: tl.constexpr):
-    if DOT:   # tensor cores: [M_PAD >= 16, BK] x [BK, BN], bf16 operands (e4m3 -> bf16 is exact)
-        acc = tl.zeros([M_PAD, BLOCK_N], dtype=tl.float32)
-        for k0 in range(0, K, BLOCK_K):
-            rk = k0 + tl.arange(0, BLOCK_K)
-            km = rk < K
-            x = tl.load(x_ptr + rm[:, None] * x_row + rk[None, :], mask=mm[:, None] & km[None, :], other=0.0)
-            w = tl.load(w_ptr + rn[:, None] * w_row + rk[None, :], mask=nm[:, None] & km[None, :], other=0.0)
-            acc = tl.dot(x.to(tl.bfloat16), tl.trans(w.to(tl.bfloat16)), acc)
-        return acc
-    acc3 = tl.zeros([M_PAD, BLOCK_N, BLOCK_K], dtype=tl.float32)
-    for k0 in range(0, K, BLOCK_K):
-        rk = k0 + tl.arange(0, BLOCK_K)
-        km = rk < K
-        x = tl.load(x_ptr + rm[:, None] * x_row + rk[None, :], mask=mm[:, None] & km[None, :], other=0.0)
-        w = tl.load(w_ptr + rn[:, None] * w_row + rk[None, :], mask=nm[:, None] & km[None, :], other=0.0)
-        acc3 += x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :]
-    return tl.sum(acc3, 2)
-
-
-@triton.jit
-def _gemv_kernel(x_ptr, x_row, w_ptr, w_row, ws_ptr, y_ptr, y_row, M, N, K, up_off, cap,
-                 M_PAD: tl.constexpr, EPI: tl.constexpr, FP8: tl.constexpr,
-                 BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, DOT: tl.constexpr = False,
-                 PDL: tl.constexpr = False):
-    """y[m, n] = x[m] . W[n] for M <= M_PAD rows. EPI 0: store bf16; 1: y += bf16(.)
-    (residual add in place); 2: SwiGLU, n < up_off reads gate row n and up row up_off+n;
-    3: fp32 logits cap * tanh(bf16(.) / cap) (cap 0: none). FP8: W e4m3 times ws[n]."""
-    _pdl(PDL)
-    pid = tl.program_id(0)
-    rn = pid * BLOCK_N + tl.arange(0, BLOCK_N)
-    nm = rn < N
-    rm = tl.arange(0, M_PAD)
-    mm = rm < M
-    acc = _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn, nm, K, M_PAD, BLOCK_N, BLOCK_K, DOT)
-    if FP8:
-        acc = acc * tl.load(ws_ptr + rn, mask=nm, other=0.0)[None, :]
-    ym = mm[:, None] & nm[None, :]
-    yp = y_ptr + rm[:, None] * y_row + rn[None, :]
-    if EPI == 2:
-        up = _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn + up_off, nm, K, M_PAD, BLOCK_N, BLOCK_K, DOT)
-        if FP8:
-            up = up * tl.load(ws_ptr + up_off + rn, mask=nm, other=0.0)[None, :]
-        gate = acc.to(tl.bfloat16).to(tl.float32)
-        up = up.to(tl.bfloat16).to(tl.float32)
-        tl.store(yp, (gate * tl.sigmoid(gate) * up).to(y_ptr.dtype.element_ty), mask=ym)
-    elif EPI == 1:
-        res = tl.load(yp, mask=ym, other=0.0).to(tl.float32)
-        tl.store(yp, (acc.to(tl.bfloat16).to(tl.float32) + res).to(y_ptr.dtype.element_ty), mask=ym)
-    elif EPI == 3:
-        z = acc.to(tl.bfloat16).to(tl.float32)
-        if cap > 0:
-            t = z / cap
-            e = tl.exp(-2.0 * tl.abs(t))
-            th = (1.0 - e) / (1.0 + e)
-            z = cap * tl.where(t < 0, -th, th)
-        tl.store(yp, z, mask=ym)
-    else:
-        tl.store(yp, acc.to(y_ptr.dtype.element_ty), mask=ym)
-
-
-def gemv(x: torch.Tensor, w: torch.Tensor, y: torch.Tensor, epi: int = 0, scale: torch.Tensor | None = None,
-         n: int | None = None, up_off: int = 0, cap: float = 0.0, config: tuple | None = None) -> None:
-    """``x`` ``(M, K)`` (row stride free), ``w`` ``(N_w, K)``, ``y`` ``(M, n)``; ``n`` output
-    columns (SwiGLU: the hidden size, ``up_off`` = its up rows' offset)."""
-    M, K = x.shape
-    N = n if n is not None else w.shape[0]
-    m_pad = max(1, triton.next_power_of_2(M))
-    fp8 = scale is not None
-    cfg = tuple(config or gemv_config(m_pad, N, K, fp8, epi))
-    bn, bk, warps, stages, dot = cfg + (0,) * (5 - len(cfg))
-    _gemv_kernel[(triton.cdiv(N, bn),)](x, x.stride(0), w, w.stride(0), scale if fp8 else w, y, y.stride(0),
-                                        M, N, K, up_off, float(cap), M_PAD=max(16, m_pad) if dot else m_pad,
-                                        EPI=epi, FP8=fp8, BLOCK_N=bn, BLOCK_K=bk, DOT=bool(dot),
-                                        num_warps=warps, num_stages=stages, PDL=PDL, launch_pdl=PDL)
-
-
-# CCGQA / QK-norm + RoPE + cache write
-
-#: Slots of the CCGQA conv-latent / value-shift rings. A step reads the two tokens before
-#: its first one, so a verify of up to RING - 3 drafted tokens never overwrites them.
-RING: tl.constexpr = tl.constexpr(32)
-
-
-@triton.jit
-def _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, cidx, LAG: tl.constexpr):
-    """The pre-conv latent of the row's token ``LAG`` back (fp32; zero across a document start):
-    the call's own qkv rows while ``t >= LAG``, else the latent ring (slot mod ``RING``)."""
-    if LAG == 0:
-        v = tl.load(qkv_ptr + r * qkv_row + col + cidx).to(tl.float32)
-    elif t >= LAG:
-        v = tl.load(qkv_ptr + (r - LAG) * qkv_row + col + cidx, mask=pos >= LAG, other=0.0).to(tl.float32)
-    else:
-        v = tl.load(lat_ptr + b * lat_b + ((slot - LAG) % RING) * lat_l + col + cidx, mask=pos >= LAG,
-                    other=0.0).to(tl.float32)
-    return tl.where(pos >= LAG, v, 0.0)
-
-
-@triton.jit
-def _conv(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cs_ptr, csb_ptr, ch_ptr, chb_ptr, h, r, t, b, slot, pos, col,
-          HD: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr):
-    """Virtual head ``h``'s causal conv (depthwise ``KS`` taps, then per-head dense ``KC``
-    taps; lag-ordered): every load is independent, one ``[HD, HD]`` tile per dense tap."""
-    d = tl.arange(0, HD)
-    b1 = tl.load(csb_ptr + h * HD + d).to(tl.float32)
-    acc = tl.load(chb_ptr + h * HD + d).to(tl.float32)
-    for j in tl.static_range(KC):
-        u = b1
-        for i in tl.static_range(KS):
-            a = tl.load(cs_ptr + (h * HD + d) * KS + i).to(tl.float32)
-            u += a * _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, d, i + j)
-        w2 = tl.load(ch_ptr + ((h * KC + j) * HD + d[:, None]) * HD + d[None, :]).to(tl.float32)
-        acc += tl.sum(w2 * u[:, None], 0)
-    return acc
-
-
-@triton.jit
-def _qk_decode_kernel(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, vh_ptr, vh_b, vh_l,
-                      pos_ptr, slot_ptr, cos_ptr, sin_ptr, cs_ptr, csb_ptr, ch_ptr, chb_ptr, qg_ptr, kg_ptr,
-                      q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h, T, eps,
-                      NQ: tl.constexpr, NKV: tl.constexpr, NXQ: tl.constexpr, HD: tl.constexpr,
-                      V_OFF: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr, CCGQA: tl.constexpr,
-                      QK_MEAN: tl.constexpr, VSHIFT: tl.constexpr, ROPE: tl.constexpr,
-                      PDL: tl.constexpr = False):
-    """One (row, virtual head): heads ``< NQ`` are queries (written to ``q``), the
-    rest keys (written to the K cache, with the value head to the V cache)."""
-    _pdl(PDL)
-    r = tl.program_id(0).to(tl.int64)
-    h = tl.program_id(1)
-    b = r // T
-    t = r % T
-    pos = tl.load(pos_ptr + r).to(tl.int64)
-    slot = tl.load(slot_ptr + r).to(tl.int64)
-    HALF: tl.constexpr = HD // 2
-    G: tl.constexpr = NQ // NKV
-    d = tl.arange(0, HD)
-    is_k = h >= NQ
-    g = h - NQ
-    ph = tl.where(is_k, NXQ + g, h % NXQ)       # physical latent head (VHA: h % NXQ)
-    col = ph * HD
-    base = qkv_ptr + r * qkv_row
-    keep = t + RING >= T   # only a call's last RING tokens go to the rings (no write races)
-    if CCGQA:
-        x = _conv(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cs_ptr, csb_ptr, ch_ptr, chb_ptr, h, r, t, b, slot,
-                  pos, col, HD, KS, KC)
-        if QK_MEAN:   # mu_q = (x_q + x_k(group)) / 2, mu_k = the group mean of mu_q
-            if is_k:
-                qm = tl.zeros([HD], dtype=tl.float32)
-                for jh in tl.static_range(G):
-                    qm += tl.load(base + ((g * G + jh) % NXQ) * HD + d).to(tl.float32)
-                x += 0.5 * (qm / G + tl.load(base + col + d).to(tl.float32))
-            else:
-                x += 0.5 * (tl.load(base + col + d).to(tl.float32)
-                            + tl.load(base + (NXQ + h // G) * HD + d).to(tl.float32))
-    else:
-        x = tl.load(base + col + d).to(tl.float32)
-    gain = tl.load(tl.where(is_k, kg_ptr + g * HD, qg_ptr) + d)
-    x = x * tl.rsqrt(tl.sum(x * x, 0) / HD + eps) * gain
-    if not CCGQA:   # plain QK-norm rounds to bf16 before RoPE
-        x = x.to(tl.bfloat16).to(tl.float32)
-    lo, hi = tl.split(tl.permute(tl.reshape(x, [2, HALF]), (1, 0)))
-    if ROPE:
-        i2 = tl.arange(0, HALF)
-        c = tl.load(cos_ptr + pos * HALF + i2)
-        s = tl.load(sin_ptr + pos * HALF + i2)
-        lo, hi = lo * c - hi * s, hi * c + lo * s
-    lo_i = tl.arange(0, HALF)
-    if is_k:
-        dst = kc_ptr + b * kc_b + g * kc_h + slot * HD
-    else:
-        dst = q_ptr + r * q_row + h * HD
-    tl.store(dst + lo_i, lo.to(q_ptr.dtype.element_ty))
-    tl.store(dst + HALF + lo_i, hi.to(q_ptr.dtype.element_ty))
-    if is_k:   # the value head: half of the flat value channels are the previous token's
-        ch = g * HD + d
-        own = tl.load(base + V_OFF + ch)
-        v = own
-        if VSHIFT:
-            VH: tl.constexpr = NKV * HD // 2
-            upper = ch >= VH
-            if t >= 1:
-                prev = tl.load(base - qkv_row + V_OFF + ch, mask=upper, other=0.0)
-            else:
-                prev = tl.load(vh_ptr + b * vh_b + ((slot - 1) % RING) * vh_l + ch - VH,
-                               mask=upper & (slot >= 1), other=0.0)
-            v = tl.where(upper, tl.where(pos > 0, prev, 0.0), own)
-            tl.store(vh_ptr + b * vh_b + (slot % RING) * vh_l + ch - VH, own, mask=upper & keep)
-        tl.store(vc_ptr + b * kc_b + g * kc_h + slot * HD + d, v.to(vc_ptr.dtype.element_ty))
-    if CCGQA:   # the latent ring: physical q head p written by query program p, k heads by theirs
-        if ((h < NXQ) | is_k) & keep:
-            tl.store(lat_ptr + b * lat_b + (slot % RING) * lat_l + col + d, tl.load(base + col + d))
-
-
-def qk_decode(qkv, lat, vh, pos, slot, cos, sin, conv, qg, kg, q, kc, vc, T: int, eps: float, *, nq, nkv,
-              nxq, hd, v_off, ks, kc_taps, ccgqa, qk_mean, vshift, rope) -> None:
-    R = qkv.shape[0]
-    cs, csb, ch, chb = conv if conv is not None else (qkv,) * 4
-    _qk_decode_kernel[(R, nq + nkv)](
-        qkv, qkv.stride(0), lat if lat is not None else qkv, lat.stride(0) if lat is not None else 0,
-        lat.stride(1) if lat is not None else 0, vh if vh is not None else qkv,
-        vh.stride(0) if vh is not None else 0, vh.stride(1) if vh is not None else 0,
-        pos, slot, cos, sin, cs, csb, ch, chb, qg, kg, q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1),
-        T, eps, NQ=nq, NKV=nkv, NXQ=nxq, HD=hd, V_OFF=v_off, KS=ks, KC=kc_taps, CCGQA=ccgqa,
-        QK_MEAN=qk_mean, VSHIFT=vshift, ROPE=rope, num_warps=8, PDL=PDL, launch_pdl=PDL)
-
-
-@triton.jit
-def _tile_lat(qkv_ptr, qkv_row, r0, t, tm, pos, col, d, T, LAG: tl.constexpr):
-    """``[BT, HD]`` pre-conv latents ``LAG`` tokens back of a prompt tile (zero across a
-    document start or before the prompt)."""
-    ok = tm & (pos >= LAG) & (t >= LAG)
-    return tl.load(qkv_ptr + (r0 + t - LAG)[:, None] * qkv_row + col + d[None, :], mask=ok[:, None],
-                   other=0.0).to(tl.float32)
-
-
-@triton.jit
-def _qk_prefill_kernel(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, vh_ptr, vh_b, vh_l, pos_ptr, cos_ptr, sin_ptr,
-                       cs_ptr, csb_ptr, ch_ptr, chb_ptr, qg_ptr, kg_ptr, q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h,
-                       T, eps, NQ: tl.constexpr, NKV: tl.constexpr, NXQ: tl.constexpr, HD: tl.constexpr,
-                       V_OFF: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr, QK_MEAN: tl.constexpr,
-                       VSHIFT: tl.constexpr, ROPE: tl.constexpr, BT: tl.constexpr):
-    """``qk_decode`` for a prompt from slot 0, ``BT`` tokens of one virtual head per program:
-    the head's dense conv taps are read once per tile and applied on tensor cores (bf16 taps,
-    the input split into bf16 hi + lo parts), not once per token."""
-    tt = tl.program_id(0)
-    h = tl.program_id(1)
-    b = tl.program_id(2).to(tl.int64)
-    HALF: tl.constexpr = HD // 2
-    G: tl.constexpr = NQ // NKV
-    t = tt * BT + tl.arange(0, BT)
-    tm = t < T
-    r0 = b * T
-    rows = r0 + t
-    pos = tl.load(pos_ptr + rows, mask=tm, other=0).to(tl.int64)
-    d = tl.arange(0, HD)
-    is_k = h >= NQ
-    g = h - NQ
-    ph = tl.where(is_k, NXQ + g, h % NXQ)
-    col = ph * HD
-    b1 = tl.load(csb_ptr + h * HD + d).to(tl.float32)
-    x = tl.zeros([BT, HD], dtype=tl.float32) + tl.load(chb_ptr + h * HD + d).to(tl.float32)[None, :]
-    for j in tl.static_range(KC):
-        u = tl.zeros([BT, HD], dtype=tl.float32) + b1[None, :]
-        for i in tl.static_range(KS):
-            a = tl.load(cs_ptr + (h * HD + d) * KS + i).to(tl.float32)
-            u += a[None, :] * _tile_lat(qkv_ptr, qkv_row, r0, t, tm, pos, col, d, T, i + j)
-        w2 = tl.load(ch_ptr + ((h * KC + j) * HD + d[:, None]) * HD + d[None, :]).to(tl.bfloat16)
-        u_hi = u.to(tl.bfloat16)   # the taps are bf16 already: u = hi + lo, two bf16 MMAs (~16-bit u)
-        x = tl.dot(u_hi, w2, x)
-        x = tl.dot((u - u_hi.to(tl.float32)).to(tl.bfloat16), w2, x)
-    base = qkv_ptr + rows[:, None] * qkv_row
-    if QK_MEAN:
-        if is_k:
-            qm = tl.zeros([BT, HD], dtype=tl.float32)
-            for jh in tl.static_range(G):
-                qm += tl.load(base + ((g * G + jh) % NXQ) * HD + d[None, :], mask=tm[:, None], other=0.0).to(
-                    tl.float32)
-            x += 0.5 * (qm / G + tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0).to(tl.float32))
-        else:
-            x += 0.5 * (tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0).to(tl.float32)
-                        + tl.load(base + (NXQ + h // G) * HD + d[None, :], mask=tm[:, None], other=0.0).to(
-                            tl.float32))
-    gain = tl.load(tl.where(is_k, kg_ptr + g * HD, qg_ptr) + d)
-    x = x * tl.rsqrt(tl.sum(x * x, 1) / HD + eps)[:, None] * gain[None, :]
-    lo, hi = tl.split(tl.permute(tl.reshape(x, [BT, 2, HALF]), (0, 2, 1)))
-    i2 = tl.arange(0, HALF)
-    if ROPE:
-        c = tl.load(cos_ptr + pos[:, None] * HALF + i2[None, :])
-        sn = tl.load(sin_ptr + pos[:, None] * HALF + i2[None, :])
-        lo, hi = lo * c - hi * sn, hi * c + lo * sn
-    if is_k:
-        dst = kc_ptr + b * kc_b + g * kc_h + t[:, None] * HD
-    else:
-        dst = q_ptr + rows[:, None] * q_row + h * HD
-    tl.store(dst + i2[None, :], lo.to(q_ptr.dtype.element_ty), mask=tm[:, None])
-    tl.store(dst + HALF + i2[None, :], hi.to(q_ptr.dtype.element_ty), mask=tm[:, None])
-    keep = tm & (t + RING >= T)
-    if is_k:
-        ch = g * HD + d
-        own = tl.load(base + V_OFF + ch[None, :], mask=tm[:, None], other=0.0)
-        v = own
-        if VSHIFT:
-            VH: tl.constexpr = NKV * HD // 2
-            upper = (ch >= VH)[None, :]
-            prev = tl.load(base - qkv_row + V_OFF + ch[None, :], mask=tm[:, None] & upper & (t >= 1)[:, None],
-                           other=0.0)
-            v = tl.where(upper, tl.where((pos > 0)[:, None], prev, 0.0), own)
-            tl.store(vh_ptr + b * vh_b + (t % RING)[:, None] * vh_l + ch[None, :] - VH, own,
-                     mask=keep[:, None] & upper)
-        tl.store(vc_ptr + b * kc_b + g * kc_h + t[:, None] * HD + d[None, :], v.to(vc_ptr.dtype.element_ty),
-                 mask=tm[:, None])
-    if (h < NXQ) | is_k:
-        tl.store(lat_ptr + b * lat_b + (t % RING)[:, None] * lat_l + col + d[None, :],
-                 tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0), mask=keep[:, None])
-
-
-def qk_prefill(qkv, lat, vh, pos, cos, sin, conv, qg, kg, q, kc, vc, T: int, eps: float, *, nq, nkv, nxq, hd,
-               v_off, ks, kc_taps, qk_mean, vshift, rope) -> None:
-    """CCGQA ``qk_decode`` for ``B`` prompts of ``T`` tokens starting at slot 0."""
-    B = qkv.shape[0] // T
-    cs, csb, ch, chb = conv
-    bt = 32 if hd <= 128 else 16   # (64-token tiles: ptxas falls to 40 registers and spills ~20x slower)
-    _qk_prefill_kernel[(triton.cdiv(T, bt), nq + nkv, B)](
-        qkv, qkv.stride(0), lat, lat.stride(0), lat.stride(1), vh if vh is not None else qkv,
-        vh.stride(0) if vh is not None else 0, vh.stride(1) if vh is not None else 0, pos, cos, sin,
-        cs, csb, ch, chb, qg, kg, q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1), T, eps, NQ=nq, NKV=nkv,
-        NXQ=nxq, HD=hd, V_OFF=v_off, KS=ks, KC=kc_taps, QK_MEAN=qk_mean, VSHIFT=vshift, ROPE=rope, BT=bt,
-        num_warps=4)
-
-
-# flash decoding
-#
-# grid (span, KV head, sequence). A program reads its span of one KV head's keys once for
-# every query of the sequence's T rows (T > 1: a speculative verify; each row has its own
-# causal limit) and all G query heads of the group, stores the unnormalised partials, and
-# the last span program of (sequence, KV head) merges them and finishes the rows.
-
-
-@triton.jit
-def _spans(start, slot, NSPL: tl.constexpr, BL: tl.constexpr):
-    """``(span, n_spans)``: keys ``[start, slot]`` cut into at most ``NSPL`` spans of a multiple
-    of ``BL`` keys, sized from the length on the device (one graph serves every length)."""
-    n = slot + 1 - start
-    span = tl.maximum(tl.cdiv(tl.cdiv(n, NSPL), BL), 1) * BL
-    return span, tl.where(n > 0, tl.cdiv(n, span), 0)
-
-
-@triton.jit
-def _merge(ml_ptr, acc_ptr, base0, n_sp, hrow, HD: tl.constexpr, G_PAD: tl.constexpr, GH_PAD: tl.constexpr,
-           NSPL: tl.constexpr, CH: tl.constexpr):
-    """Merge the ``n_sp`` span partials of the query heads ``hrow`` (in the group): the maxima
-    in one load, the accumulators ``CH`` spans at a time -> bf16-rounded output."""
-    d = tl.arange(0, HD)
-    si = tl.arange(0, NSPL)
-    m = tl.load(ml_ptr + (base0 + si)[:, None] * 2 * G_PAD + hrow[None, :], mask=(si < n_sp)[:, None],
-                other=float("-inf"))
-    l = tl.load(ml_ptr + (base0 + si)[:, None] * 2 * G_PAD + G_PAD + hrow[None, :], mask=(si < n_sp)[:, None],
-                other=0.0)
-    mg = tl.max(m, 0)
-    mg = tl.where(mg == float("-inf"), 0.0, mg)
-    lsum = tl.sum(tl.exp2(m - mg[None, :]) * l, 0)
-    o = tl.zeros([GH_PAD, HD], dtype=tl.float32)
-    for c0 in range(0, n_sp, CH):
-        cs = c0 + tl.arange(0, CH)
-        cm = cs < n_sp
-        mc = tl.load(ml_ptr + (base0 + cs)[:, None] * 2 * G_PAD + hrow[None, :], mask=cm[:, None],
-                     other=float("-inf"))
-        a = tl.load(acc_ptr + ((base0 + cs)[:, None, None] * G_PAD + hrow[None, :, None]) * HD
-                    + d[None, None, :], mask=cm[:, None, None], other=0.0)
-        o += tl.sum(tl.exp2(mc - mg[None, :])[:, :, None] * a, 0)
-    o = o / tl.maximum(lsum, 1e-30)[:, None]
-    return o.to(tl.bfloat16).to(tl.float32)
-
-
-@triton.jit
-def _finish(ml_ptr, acc_ptr, r, g, b, n_sp, slot, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row,
-            lam_off, gate_off, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr,
-            HD: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr, GATE: tl.constexpr):
-    """Merge row ``r``'s spans for KV head ``g`` and finish DiffAttn V2 + XSA (+ gate) into ``y``."""
-    base0 = (r * NKV + g) * NSPL
-    d = tl.arange(0, HD)
-    if DIFF:
-        GH: tl.constexpr = G // 2
-        GH_PAD: tl.constexpr = G_PAD // 2
-        i = tl.arange(0, GH_PAD)
-        im = i < GH
-        y0 = _merge(ml_ptr, acc_ptr, base0, n_sp, 2 * i, HD, G_PAD, GH_PAD, NSPL, CH)
-        y1 = _merge(ml_ptr, acc_ptr, base0, n_sp, 2 * i + 1, HD, G_PAD, GH_PAD, NSPL, CH)
-        lam = tl.load(qkv_ptr + r * qkv_row + lam_off + g * GH + i, mask=im, other=0.0).to(tl.float32)
-        y = y0 - tl.sigmoid(lam)[:, None] * y1
-    else:
-        GH: tl.constexpr = G
-        GH_PAD: tl.constexpr = G_PAD
-        i = tl.arange(0, GH_PAD)
-        im = i < GH
-        y = _merge(ml_ptr, acc_ptr, base0, n_sp, i, HD, G_PAD, GH_PAD, NSPL, CH)
-    if XSA:
-        v = tl.load(vc_ptr + b * kc_b + g * kc_h + slot * HD + d).to(tl.float32)
-        vn = v / tl.maximum(tl.sqrt(tl.sum(v * v, 0)), 1e-12)
-        y = y - tl.sum(y * vn[None, :], 1)[:, None] * vn[None, :]
-    oh = g * GH + i
-    if GATE == 1:   # per channel
-        gt = tl.load(qkv_ptr + r * qkv_row + gate_off + oh[:, None] * HD + d[None, :], mask=im[:, None],
-                     other=0.0).to(tl.float32)
-        y = y * tl.sigmoid(gt)
-    elif GATE == 2:   # per head
-        gt = tl.load(qkv_ptr + r * qkv_row + gate_off + oh, mask=im, other=0.0).to(tl.float32)
-        y = y * tl.sigmoid(gt)[:, None]
-    tl.store(y_ptr + r * y_row + oh[:, None] * HD + d[None, :], y.to(y_ptr.dtype.element_ty), mask=im[:, None])
-
-
-@triton.jit
-def _fence():
-    """``fence.acq_rel.gpu`` in every thread (CUDA's ``__threadfence``)."""
-    tl.inline_asm_elementwise("fence.acq_rel.gpu; mov.u32 $0, 0;", "=r", [], dtype=tl.int32, is_pure=False,
-                              pack=1)
-
-
-@triton.jit
-def _attn_kernel(q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h, slot_ptr, start_ptr, ml_ptr, acc_ptr, cnt_ptr,
-                 qkv_ptr, qkv_row, y_ptr, y_row, T, qk_scale, lam_off, gate_off,
-                 NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr, QR: tl.constexpr,
-                 HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr,
-                 GATE: tl.constexpr, FINISH: tl.constexpr, PDL: tl.constexpr = False):
-    """``FINISH`` (one row per sequence): the last span program finishes the row; else
-    ``_finish_kernel`` does, a program per row."""
-    _pdl(PDL)
-    sp = tl.program_id(0)
-    g = tl.program_id(1)
-    b = tl.program_id(2).to(tl.int64)
-    r0 = b * T
-    start = tl.load(start_ptr + r0)
-    last = tl.load(slot_ptr + r0 + T - 1)            # the sequence's newest row reaches furthest
-    span, n_sp = _spans(start, last, NSPL, BL)
-    lo = start + sp * span
-    hi = tl.minimum(last + 1, lo + span)
-    if lo < hi:
-        qi = tl.arange(0, QR)                         # query row = (token t, head j of the group)
-        t = qi // G
-        j = qi % G
-        qm = t < T
-        lim = tl.load(slot_ptr + r0 + t, mask=qm, other=-1)   # each row's own causal limit
-        d = tl.arange(0, HD)
-        q = tl.load(q_ptr + (r0 + t)[:, None] * q_row + (g * G + j)[:, None] * HD + d[None, :],
-                    mask=qm[:, None], other=0.0)
-        m_i = tl.full([QR], float("-inf"), tl.float32)
-        l_i = tl.zeros([QR], dtype=tl.float32)
-        acc = tl.zeros([QR, HD], dtype=tl.float32)
-        kbase = b * kc_b + g * kc_h
-        for k0 in range(lo, hi, BL):
-            kk = k0 + tl.arange(0, BL)
-            km = kk < hi
-            k = tl.load(kc_ptr + kbase + kk[:, None] * HD + d[None, :], mask=km[:, None], other=0.0)
-            s = tl.dot(q, tl.trans(k)) * qk_scale                       # base-2 logits
-            s = tl.where(km[None, :] & (kk[None, :] <= lim[:, None]), s, float("-inf"))
-            m_new = tl.maximum(m_i, tl.max(s, 1))
-            m_ok = tl.where(m_new == float("-inf"), 0.0, m_new)         # a row with no key yet
-            alpha = tl.exp2(m_i - m_ok)
-            p = tl.exp2(s - m_ok[:, None])
-            l_i = l_i * alpha + tl.sum(p, 1)
-            v = tl.load(vc_ptr + kbase + kk[:, None] * HD + d[None, :], mask=km[:, None], other=0.0)
-            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
-            m_i = m_new
-        base = ((r0 + t) * NKV + g) * NSPL + sp
-        tl.store(ml_ptr + base * 2 * G_PAD + j, m_i, mask=qm)
-        tl.store(ml_ptr + base * 2 * G_PAD + G_PAD + j, l_i, mask=qm)
-        tl.store(acc_ptr + (base * G_PAD + j)[:, None] * HD + d[None, :], acc, mask=qm[:, None])
-        if FINISH:
-            _fence()
-            tl.debug_barrier()
-            old = tl.atomic_add(cnt_ptr + b * NKV + g, 1, sem="acq_rel", scope="gpu")
-            if old == n_sp - 1:   # the last span merges and finishes the row
-                _fence()
-                tl.atomic_xchg(cnt_ptr + b * NKV + g, 0, sem="relaxed", scope="gpu")
-                _finish(ml_ptr, acc_ptr, r0, g, b, n_sp, last, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row,
-                        lam_off, gate_off, NKV, NSPL, G, G_PAD, HD, CH, DIFF, XSA, GATE)
-
-
-@triton.jit
-def _finish_kernel(ml_ptr, slot_ptr, start_ptr, acc_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row, T,
-                   lam_off, gate_off, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr,
-                   HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr,
-                   GATE: tl.constexpr, PDL: tl.constexpr = False):
-    """The finish of a multi-row (verify / short prompt) attention, a program per (row, KV head)."""
-    _pdl(PDL)
-    r = tl.program_id(0).to(tl.int64)
-    g = tl.program_id(1)
-    b = r // T
-    _, n_sp = _spans(tl.load(start_ptr + b * T), tl.load(slot_ptr + b * T + T - 1), NSPL, BL)
-    _finish(ml_ptr, acc_ptr, r, g, b, n_sp, tl.load(slot_ptr + r), qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
-            y_row, lam_off, gate_off, NKV, NSPL, G, G_PAD, HD, CH, DIFF, XSA, GATE)
-
-
-#: Smallest query tile of the decode attention, and whether a 1-row step finishes in the
-#: attention kernel (saves a launch) or in ``_finish_kernel`` like a verify. A row's sums
-#: round the same way only in the same tile and code: tests set 32 / False so a 1-row step
-#: and a verify match bitwise.
-QR_MIN = 16
-FUSED_FINISH = True
-
-
-def attention(q, kc, vc, slot, start, qkv, y, ml, acc, cnt, T: int, *, nq, nkv, hd, nspl, lam_off,
-              gate_off, diff, xsa, gate) -> None:
-    """``q`` ``(R, nq, hd)`` for ``R = B * T`` rows (``T`` consecutive tokens per sequence),
-    caches ``(B, nkv, L, hd)``, ``y`` ``(R, nh * hd)``; scratch ``ml`` ``(R, nkv, nspl, 2, G_PAD)``
-    / ``acc`` ``(R, nkv, nspl, G_PAD, hd)`` fp32 and zeroed counters ``cnt`` (>= ``B * nkv``
-    int32, left zeroed)."""
-    R = q.shape[0]
-    G = nq // nkv
-    g_pad = max(16, triton.next_power_of_2(G))
-    qr = max(QR_MIN, triton.next_power_of_2(T * G))
-    ch = 8 if g_pad <= 16 else 4
-    _attn_kernel[(nspl, nkv, R // T)](q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1), slot, start, ml, acc,
-                                      cnt, qkv, qkv.stride(0), y, y.stride(0), T, hd ** -0.5 * LOG2E, lam_off,
-                                      gate_off,
-                                      NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad, QR=qr, HD=hd, BL=64,
-                                      CH=ch, DIFF=diff, XSA=xsa, GATE=gate, FINISH=T == 1 and FUSED_FINISH,
-                                      num_warps=4 if qr <= 32 else 8, num_stages=2, PDL=PDL, launch_pdl=PDL)
-    if T > 1 or not FUSED_FINISH:
-        _finish_kernel[(R, nkv)](ml, slot, start, acc, qkv, qkv.stride(0), vc, kc.stride(0), kc.stride(1), y,
-                                 y.stride(0), T, lam_off, gate_off, NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad, HD=hd,
-                                 BL=64, CH=ch, DIFF=diff, XSA=xsa, GATE=gate, num_warps=4, PDL=PDL,
-                                 launch_pdl=PDL)
-
-
-# sampling
-
-
-@triton.jit
-def _sample_kernel(x_ptr, x_row, V, inv_temp, top_k, top_p, min_p, seed_ptr, excl_ptr, y_ptr, pd_ptr, u_ptr,
-                   pres_ptr, pres_row, ids_ptr, penalty, T,
-                   TOP_K: tl.constexpr, TOP_P: tl.constexpr, MIN_P: tl.constexpr, EXCL: tl.constexpr,
-                   REP: tl.constexpr, GREEDY: tl.constexpr, T_MAX: tl.constexpr, BLOCK: tl.constexpr,
-                   ITERS: tl.constexpr):
-    """One row: transformers' repetition penalty -> temperature -> top-k -> top-p -> min-p,
-    then a draw (``GREEDY``: the argmax of the penalised logits).
-
-    ``REP``: tokens in the row's history (``pres[b]``, rows ``b * T + t``) and in this call's
-    own inputs ``ids[b, 1..t]`` (a verify's earlier draft tokens) get ``x < 0 ? x * penalty :
-    x / penalty``.
-
-    Thresholds instead of a sort: top-k keeps ``x >= `` the k-th largest value, top-p
-    keeps ``p >= tau`` for the largest ``tau`` whose kept mass reaches ``top_p`` (the
-    smallest top set, as transformers' ascending cumsum), both found by bisection to
-    below an ulp. The draw inverts the CDF of the kept mass with one Philox uniform.
-    ``EXCL``: also the probability of token ``excl`` under the result (``pd``, for
-    speculative acceptance, with a uniform ``u``) and the draw made without it."""
-    r = tl.program_id(0).to(tl.int64)
-    offs = tl.arange(0, BLOCK)
-    msk = offs < V
-    x = tl.load(x_ptr + r * x_row + offs, mask=msk, other=float("-inf"))
-    if REP:
-        b = r // T
-        t = r % T
-        seen = tl.load(pres_ptr + b * pres_row + offs, mask=msk, other=0) != 0
-        for i in tl.static_range(1, T_MAX):
-            if i <= t:
-                seen = seen | (offs == tl.load(ids_ptr + b * T + i))
-        x = tl.where(seen, tl.where(x < 0, x * penalty, x / penalty), x)
-    if GREEDY:
-        tl.store(y_ptr + r, tl.argmax(x, 0).to(tl.int64))
-    else:
-        x = x * inv_temp
-        if TOP_K:
-            lo = tl.min(tl.where(msk, x, float("inf")), 0)
-            hi = tl.max(x, 0)
-            for _ in range(ITERS):   # count(x >= lo) >= k > count(x >= hi)
-                mid = lo + (hi - lo) * 0.5
-                ok = tl.sum((x >= mid).to(tl.int32), 0) >= top_k
-                lo = tl.where(ok, mid, lo)
-                hi = tl.where(ok, hi, mid)
-            kth = tl.min(tl.where(x >= lo, x, float("inf")), 0)
-            x = tl.where(x >= kth, x, float("-inf"))
-        m = tl.max(x, 0)
-        p = tl.exp(x - m)                      # max 1
-        z = tl.sum(p, 0)
-        if TOP_P:
-            lo = tl.zeros([], dtype=tl.float32)
-            hi = tl.full([], 1.0, tl.float32)  # mass(p >= lo) >= top_p * z > mass(p >= hi) (or hi = 1)
-            goal = top_p * z
-            for _ in range(ITERS):
-                mid = lo + (hi - lo) * 0.5
-                ok = tl.sum(tl.where(p >= mid, p, 0.0), 0) >= goal
-                lo = tl.where(ok, mid, lo)
-                hi = tl.where(ok, hi, mid)
-            tau = tl.min(tl.where(p >= lo, p, float("inf")), 0)
-            p = tl.where(p >= tau, p, 0.0)
-        if MIN_P:
-            p = tl.where(p >= min_p, p, 0.0)   # max is 1: min_p * max = min_p
-        seed = tl.load(seed_ptr)
-        if EXCL:
-            e = tl.load(excl_ptr + r)
-            pd = tl.sum(tl.where(offs == e, p, 0.0), 0) / tl.sum(p, 0)
-            tl.store(pd_ptr + r, pd)
-            tl.store(u_ptr + r, tl.rand(seed + 7919, r.to(tl.int32)))
-            p = tl.where(offs == e, 0.0, p)
-        # inverse CDF: one uniform per row; the first kept token whose running mass passes it
-        cdf = tl.cumsum(p, 0)
-        u = tl.rand(seed, r.to(tl.int32)) * tl.max(cdf, 0)
-        tok = tl.min(tl.where((cdf > u) & (p > 0), offs, BLOCK), 0)
-        tl.store(y_ptr + r, tl.where(tok < V, tok, tl.argmax(p, 0)).to(tl.int64))
-
-
-def sample(x: torch.Tensor, seed: torch.Tensor, y: torch.Tensor, *, temperature: float, top_k: int, top_p: float,
-           min_p: float, excl: torch.Tensor | None = None, pd: torch.Tensor | None = None,
-           u: torch.Tensor | None = None, greedy: bool = False, penalty: float = 1.0,
-           pres: torch.Tensor | None = None, ids: torch.Tensor | None = None, T: int = 1) -> None:
-    """Draw ``y`` ``(R,)`` from fp32 logits ``x`` ``(R, V)``; ``seed`` a 1-element int64 device
-    tensor (advance it between calls). With ``excl``: ``pd`` / ``u`` as in ``_sample_kernel``.
-    ``penalty != 1``: the repetition penalty over ``pres`` ``(B, V)`` bool and the call's input
-    ids ``(B * T,)`` (rows ``b * T + t``)."""
-    R, V = x.shape
-    block = triton.next_power_of_2(V)
-    rep = penalty != 1.0
-    _sample_kernel[(R,)](x, x.stride(0), V, 1.0 / temperature, int(top_k), float(top_p), float(min_p), seed,
-                         excl if excl is not None else y, y, pd if pd is not None else x,
-                         u if u is not None else x, pres.view(torch.uint8) if rep else y,
-                         pres.stride(0) if rep else 0, ids if rep else y, float(penalty), T,
-                         TOP_K=bool(top_k) and top_k < V, TOP_P=top_p < 1.0, MIN_P=bool(min_p),
-                         EXCL=excl is not None, REP=rep, GREEDY=greedy, T_MAX=T, BLOCK=block, ITERS=40,
-                         num_warps=32 if block >= 16384 else 8)
-
-
-kernels = sys.modules[__name__]   # ``from .luau_kernels import kernels as K``: transformers'
-                                  # remote code only follows ``from .module import ...`` lines
+"""Triton kernels for ``luau_infer`` (fused decoding of the LuauCoder HF port).
+
+Every kernel works on *rows*: one row per token of the call (``B`` rows when
+decoding, ``B * T`` for a prefill or a speculative verify), each with its own
+document position ``pos`` (RoPE, conv / value-shift masks), cache ``slot``
+(where its K/V go) and ``start`` (the first valid key slot of its sequence:
+left padding). Nothing reads a Python scalar that changes between steps, so a
+decode step replays as a CUDA graph.
+
+* ``mix_norm``     Block AttnRes depth mix (a softmax over the sources, one tile)
+                   + RMSNorm, or the embedding gather + RMSNorm (layer 0).
+* ``gemv``         skinny ``x @ W.T`` (bf16 or rowwise-FP8 weights), epilogues:
+                   store, residual add, SwiGLU (gate|up rows), softcap logits.
+* ``qk_decode``    CCGQA conv + QK-mean + QK-norm + gains + RoPE + value shift,
+                   writes K/V to the cache and the conv latents to their ring.
+* ``attention``    GQA flash-decoding: each program reads a span of one KV
+                   head's keys once for all its query heads and all the
+                   sequence's rows (a speculative verify); the last span
+                   program merges and finishes DiffAttn V2 + XSA (+ gate).
+* ``qk_prefill``   ``qk_decode`` for prompts, in token tiles.
+* ``iha_mix``      IHA: pseudo-head mixes of the latent q / k / v, norm, temperature,
+                   RoPE at virtual positions, interleaved pseudo-K/V cache write;
+                   ``attention`` then runs over pseudo-rows with a sliding window and
+                   its finish collapses them (+ Gram-Schmidt XSA, gate).
+* ``head_gate``    the gated byte head's ``[h, sigmoid(u.h + b) P h]`` input.
+* ``sample``       repetition penalty / temperature / top-k / top-p / min-p + draw
+                   (or greedy), one row per program.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+try:   # (remote code: transformers only requires imports outside try blocks)
+    import triton
+    import triton.language as tl
+except ImportError as exc:  # pragma: no cover
+    raise ImportError("the fused LuauCoder decoding kernels need triton") from exc
+
+from .luau_gemv_tuned import TABLE as _TABLE
+
+LOG2E = 1.4426950408889634
+
+# programmatic dependent launch (SM90+); LUAU_INFER_PDL=0 turns it off
+PDL = (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
+       and os.environ.get("LUAU_INFER_PDL", "1") != "0")
+
+# L2 prefetch of the next GEMV's weights (SM90+); LUAU_INFER_PREFETCH=0 turns it off
+BULK = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
+PREFETCH = BULK and os.environ.get("LUAU_INFER_PREFETCH", "1") != "0"
+
+
+def _pf(pf) -> dict:
+    """Launch kwargs of an optional prefetch target ``pf`` (a tensor)."""
+    if pf is None or not PREFETCH:
+        return dict(pf_ptr=None, pf_bytes=0, PF=False, BULK=False)
+    return dict(pf_ptr=pf, pf_bytes=pf.numel() * pf.element_size(), PF=True, BULK=BULK)
+
+
+@triton.jit
+def _pdl(ON: tl.constexpr):
+    """Wait for the previous kernel's writes, then let the next kernel launch."""
+    if ON:
+        tl.extra.cuda.gdc_wait()
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+@triton.jit
+def _prefetch(ptr, nbytes, PF: tl.constexpr, BULK: tl.constexpr):
+    """Fire-and-forget L2 prefetch of ``nbytes`` at ``ptr``, split over the kernel's programs."""
+    if PF:
+        pid = tl.program_id(0) + tl.num_programs(0) * (tl.program_id(1) + tl.num_programs(1) * tl.program_id(2))
+        npg = tl.num_programs(0) * tl.num_programs(1) * tl.num_programs(2)
+        CH: tl.constexpr = 4096 if BULK else 128
+        n = tl.cdiv(nbytes, CH)
+        per = tl.cdiv(n, npg)
+        lo = pid * per
+        hi = tl.minimum(n, lo + per)
+        base = ptr.to(tl.pointer_type(tl.int8))
+        for i0 in range(lo, hi, 256):
+            i = i0 + tl.arange(0, 256)
+            m = (i < hi).to(tl.int32)
+            a = base + i.to(tl.int64) * CH
+            if BULK:
+                sz = tl.minimum(nbytes - i.to(tl.int64) * CH, CH).to(tl.int32)
+                tl.inline_asm_elementwise(
+                    "{ .reg .pred p; setp.ne.b32 p, $3, 0; @p cp.async.bulk.prefetch.L2.global [$1], $2; "
+                    "mov.u32 $0, 0; }", "=r,l,r,r", [a, sz, m], dtype=tl.int32, is_pure=False, pack=1)
+            else:
+                tl.inline_asm_elementwise(
+                    "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p prefetch.global.L2 [$1]; mov.u32 $0, 0; }",
+                    "=r,l,r", [a, m], dtype=tl.int32, is_pure=False, pack=1)
+
+
+@triton.jit
+def _mix_norm_kernel(s_ptr, s_row, s_ent, ids_ptr, emb_ptr, w_ptr, g_ptr, out_ptr, o_row, eps, D,
+                     NE: tl.constexpr, NE_PAD: tl.constexpr, MODE: tl.constexpr, BLOCK: tl.constexpr,
+                     ADD: tl.constexpr = 0, PART: tl.constexpr = 1, PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
+    """MODE 0: h = S[r, 0]; 1: embedding gather; 2: AttnRes mix; 3: sum of entries. out = rmsnorm(h) * g."""
+    r = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, BLOCK)
+    msk = offs < D
+    g = tl.load(g_ptr + offs, mask=msk, other=0.0).to(tl.float32)
+    if MODE == 2:
+        w = tl.load(w_ptr + offs, mask=msk, other=0.0)
+    _prefetch(pf_ptr, pf_bytes, PF, BULK)
+    _pdl(PDL)
+    if MODE == 1:
+        tok = tl.load(ids_ptr + r).to(tl.int64)
+        hb = tl.load(emb_ptr + tok * D + offs, mask=msk, other=0.0)
+        tl.store(s_ptr + r * s_row + offs, hb, mask=msk)
+        h = hb.to(tl.float32)
+    elif MODE == 0:
+        h = tl.load(s_ptr + r * s_row + offs, mask=msk, other=0.0).to(tl.float32)
+    else:
+        j = tl.arange(0, NE_PAD)
+        jm = j < NE
+        v = tl.load(s_ptr + r * s_row + j[:, None] * s_ent + offs[None, :], mask=jm[:, None] & msk[None, :],
+                    other=0.0).to(tl.float32)
+        if MODE == 3:
+            h = tl.sum(v, 0)
+        else:
+            sc = tl.sum(v * w[None, :], 1) * tl.rsqrt(tl.sum(v * v, 1) / D + eps)
+            sc = tl.where(j < NE - 1 + PART, sc, float("-inf"))
+            e = tl.exp(sc - tl.max(sc, 0))
+            p = e / tl.sum(e, 0)
+            h = tl.sum((p + ADD)[:, None] * v, 0)
+        h = h.to(out_ptr.dtype.element_ty).to(tl.float32)   # the mix output is bf16
+    y = h * tl.rsqrt(tl.sum(h * h, 0) / D + eps) * g
+    tl.store(out_ptr + r * o_row + offs, y.to(out_ptr.dtype.element_ty), mask=msk)
+
+
+def mix_norm(S: torch.Tensor, n_ent: int, mode: int, out: torch.Tensor, g: torch.Tensor, eps: float,
+             w: torch.Tensor | None = None, ids: torch.Tensor | None = None,
+             emb: torch.Tensor | None = None, add: bool = False, partial: bool = True, pf=None) -> None:
+    """``S`` ``(R, NS, D)`` sources (row-major per row), ``out`` ``(R, D)``."""
+    R, _, D = S.shape
+    block = triton.next_power_of_2(D)
+    ne_pad = triton.next_power_of_2(n_ent) if mode >= 2 else 1
+    warps = max(4, min(16, ne_pad * block // 1024))
+    _mix_norm_kernel[(R,)](S, S.stride(0), S.stride(1), ids if ids is not None else S,
+                           emb if emb is not None else S, w if w is not None else g, g, out,
+                           out.stride(0), eps, D, NE=n_ent, NE_PAD=ne_pad, MODE=mode, BLOCK=block,
+                           ADD=int(add), PART=int(partial), num_warps=warps, PDL=PDL, launch_pdl=PDL, **_pf(pf))
+
+
+_TUNED_PATH = Path(__file__).with_name("luau_gemv_tuned.py")
+_TUNED: dict | None = None
+
+
+def _tuned() -> dict:
+    """Pinned GEMV configs for this GPU (``benchmark/tune_gemv.py``), else {}."""
+    global _TUNED
+    if _TUNED is None:
+        name = torch.cuda.get_device_name()
+        p = torch.cuda.get_device_properties(0)
+        arch = [p.major, p.minor, p.multi_processor_count]
+        hit = _TABLE.get(name) or next((e for e in _TABLE.values() if e.get("arch") == arch), None)
+        _TUNED = dict(hit["gemv"]) if hit else {}
+    return _TUNED
+
+
+def gemv_key(m_pad: int, n: int, k: int, fp8: bool, epi: int) -> str:
+    return f"{m_pad},{n},{k},{int(fp8)},{int(epi == 2)}"
+
+
+def gemv_config(m_pad: int, n: int, k: int, fp8: bool, epi: int) -> tuple[int, int, int, int, int]:
+    """``(BLOCK_N, BLOCK_K, num_warps, num_stages, DOT)``: pinned, else a default."""
+    hit = _tuned().get(gemv_key(m_pad, n, k, fp8, epi))
+    if hit:
+        return tuple(hit) + (0,) * (5 - len(hit))
+    if m_pad >= 4:
+        return 32, 128, 4, 3, 1
+    bk = 512 if fp8 else 256
+    bn = max(1, min(16, 2048 // (bk * m_pad) * 4))
+    return bn, bk, 4, 2, 0
+
+
+@triton.jit
+def _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn, nm, K, M_PAD: tl.constexpr, BLOCK_N: tl.constexpr,
+               BLOCK_K: tl.constexpr, DOT: tl.constexpr):
+    if DOT:
+        acc = tl.zeros([M_PAD, BLOCK_N], dtype=tl.float32)
+        for k0 in range(0, K, BLOCK_K):
+            rk = k0 + tl.arange(0, BLOCK_K)
+            km = rk < K
+            x = tl.load(x_ptr + rm[:, None] * x_row + rk[None, :], mask=mm[:, None] & km[None, :], other=0.0)
+            w = tl.load(w_ptr + rn[:, None] * w_row + rk[None, :], mask=nm[:, None] & km[None, :], other=0.0)
+            acc = tl.dot(x.to(tl.bfloat16), tl.trans(w.to(tl.bfloat16)), acc)
+        return acc
+    acc3 = tl.zeros([M_PAD, BLOCK_N, BLOCK_K], dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        rk = k0 + tl.arange(0, BLOCK_K)
+        km = rk < K
+        x = tl.load(x_ptr + rm[:, None] * x_row + rk[None, :], mask=mm[:, None] & km[None, :], other=0.0)
+        w = tl.load(w_ptr + rn[:, None] * w_row + rk[None, :], mask=nm[:, None] & km[None, :], other=0.0)
+        acc3 += x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :]
+    return tl.sum(acc3, 2)
+
+
+@triton.jit
+def _gemv_kernel(x_ptr, x_row, w_ptr, w_row, ws_ptr, y_ptr, y_row, M, N, K, up_off, cap,
+                 M_PAD: tl.constexpr, EPI: tl.constexpr, FP8: tl.constexpr,
+                 BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, DOT: tl.constexpr = False,
+                 PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
+    """y[m, n] = x[m] . W[n]. EPI 0: store; 1: residual add; 2: SwiGLU; 3: softcapped fp32 logits."""
+    _prefetch(pf_ptr, pf_bytes, PF, BULK)
+    _pdl(PDL)
+    pid = tl.program_id(0)
+    rn = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    nm = rn < N
+    rm = tl.arange(0, M_PAD)
+    mm = rm < M
+    acc = _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn, nm, K, M_PAD, BLOCK_N, BLOCK_K, DOT)
+    if FP8:
+        acc = acc * tl.load(ws_ptr + rn, mask=nm, other=0.0)[None, :]
+    ym = mm[:, None] & nm[None, :]
+    yp = y_ptr + rm[:, None] * y_row + rn[None, :]
+    if EPI == 2:
+        up = _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn + up_off, nm, K, M_PAD, BLOCK_N, BLOCK_K, DOT)
+        if FP8:
+            up = up * tl.load(ws_ptr + up_off + rn, mask=nm, other=0.0)[None, :]
+        gate = acc.to(tl.bfloat16).to(tl.float32)
+        up = up.to(tl.bfloat16).to(tl.float32)
+        tl.store(yp, (gate * tl.sigmoid(gate) * up).to(y_ptr.dtype.element_ty), mask=ym)
+    elif EPI == 1:
+        res = tl.load(yp, mask=ym, other=0.0).to(tl.float32)
+        tl.store(yp, (acc.to(tl.bfloat16).to(tl.float32) + res).to(y_ptr.dtype.element_ty), mask=ym)
+    elif EPI == 3:
+        z = acc.to(tl.bfloat16).to(tl.float32)
+        if cap > 0:
+            t = z / cap
+            e = tl.exp(-2.0 * tl.abs(t))
+            th = (1.0 - e) / (1.0 + e)
+            z = cap * tl.where(t < 0, -th, th)
+        tl.store(yp, z, mask=ym)
+    else:
+        tl.store(yp, acc.to(y_ptr.dtype.element_ty), mask=ym)
+
+
+def gemv(x: torch.Tensor, w: torch.Tensor, y: torch.Tensor, epi: int = 0, scale: torch.Tensor | None = None,
+         n: int | None = None, up_off: int = 0, cap: float = 0.0, config: tuple | None = None, pf=None) -> None:
+    """``x`` ``(M, K)`` (row stride free), ``w`` ``(N_w, K)``, ``y`` ``(M, n)``; ``n`` output
+    columns (SwiGLU: the hidden size, ``up_off`` = its up rows' offset)."""
+    M, K = x.shape
+    N = n if n is not None else w.shape[0]
+    m_pad = max(1, triton.next_power_of_2(M))
+    fp8 = scale is not None
+    cfg = tuple(config or gemv_config(m_pad, N, K, fp8, epi))
+    bn, bk, warps, stages, dot = cfg + (0,) * (5 - len(cfg))
+    _gemv_kernel[(triton.cdiv(N, bn),)](x, x.stride(0), w, w.stride(0), scale if fp8 else w, y, y.stride(0),
+                                        M, N, K, up_off, float(cap), M_PAD=max(16, m_pad) if dot else m_pad,
+                                        EPI=epi, FP8=fp8, BLOCK_N=bn, BLOCK_K=bk, DOT=bool(dot),
+                                        num_warps=warps, num_stages=stages, PDL=PDL, launch_pdl=PDL, **_pf(pf))
+
+
+# CCGQA conv-latent / value-shift ring slots; a verify of up to RING - 3 drafts never overwrites
+RING: tl.constexpr = tl.constexpr(32)
+
+
+@triton.jit
+def _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, cidx, LAG: tl.constexpr):
+    """The pre-conv latent of the row's token ``LAG`` back (zero across a document start)."""
+    if LAG == 0:
+        v = tl.load(qkv_ptr + r * qkv_row + col + cidx).to(tl.float32)
+    elif t >= LAG:
+        v = tl.load(qkv_ptr + (r - LAG) * qkv_row + col + cidx, mask=pos >= LAG, other=0.0).to(tl.float32)
+    else:
+        v = tl.load(lat_ptr + b * lat_b + ((slot - LAG) % RING) * lat_l + col + cidx, mask=pos >= LAG,
+                    other=0.0).to(tl.float32)
+    return tl.where(pos >= LAG, v, 0.0)
+
+
+@triton.jit
+def _conv(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cs_ptr, csb_ptr, ch_ptr, chb_ptr, h, r, t, b, slot, pos, col,
+          HD: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr):
+    """Virtual head ``h``'s causal conv."""
+    d = tl.arange(0, HD)
+    b1 = tl.load(csb_ptr + h * HD + d).to(tl.float32)
+    acc = tl.load(chb_ptr + h * HD + d).to(tl.float32)
+    for j in tl.static_range(KC):
+        u = b1
+        for i in tl.static_range(KS):
+            a = tl.load(cs_ptr + (h * HD + d) * KS + i).to(tl.float32)
+            u += a * _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, d, i + j)
+        w2 = tl.load(ch_ptr + ((h * KC + j) * HD + d[:, None]) * HD + d[None, :]).to(tl.float32)
+        acc += tl.sum(w2 * u[:, None], 0)
+    return acc
+
+
+@triton.jit
+def _conv_pre(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, b1, acc, a0, a1, w0, w1, r, t, b, slot, pos, col,
+              HD: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr):
+    """``_conv`` (``KS``, ``KC`` <= 2) on taps already in registers (loaded before the PDL wait)."""
+    d = tl.arange(0, HD)
+    for j in tl.static_range(KC):
+        u = b1
+        for i in tl.static_range(KS):
+            a = a0 if i == 0 else a1
+            u += a * _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, d, i + j)
+        w2 = w0 if j == 0 else w1
+        acc += tl.sum(w2 * u[:, None], 0)
+    return acc
+
+
+@triton.jit
+def _qk_decode_kernel(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, vh_ptr, vh_b, vh_l,
+                      pos_ptr, slot_ptr, cos_ptr, sin_ptr, cs_ptr, csb_ptr, ch_ptr, chb_ptr, qg_ptr, kg_ptr,
+                      q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h, T, eps,
+                      NQ: tl.constexpr, NKV: tl.constexpr, NXQ: tl.constexpr, HD: tl.constexpr,
+                      V_OFF: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr, CCGQA: tl.constexpr,
+                      QK_MEAN: tl.constexpr, VSHIFT: tl.constexpr, ROPE: tl.constexpr,
+                      PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
+    """One (row, virtual head): queries go to ``q``, keys / values to the caches."""
+    h = tl.program_id(1)
+    PRE: tl.constexpr = CCGQA and KS <= 2 and KC <= 2
+    if PRE:
+        dd = tl.arange(0, HD)
+        cb1 = tl.load(csb_ptr + h * HD + dd).to(tl.float32)
+        cacc = tl.load(chb_ptr + h * HD + dd).to(tl.float32)
+        ca0 = tl.load(cs_ptr + (h * HD + dd) * KS).to(tl.float32)
+        ca1 = tl.load(cs_ptr + (h * HD + dd) * KS + (KS - 1)).to(tl.float32)
+        cw0 = tl.load(ch_ptr + ((h * KC) * HD + dd[:, None]) * HD + dd[None, :]).to(tl.float32)
+        cw1 = tl.load(ch_ptr + ((h * KC + (KC - 1)) * HD + dd[:, None]) * HD + dd[None, :]).to(tl.float32)
+    r = tl.program_id(0).to(tl.int64)
+    b = r // T
+    t = r % T
+    pos = tl.load(pos_ptr + r).to(tl.int64)
+    slot = tl.load(slot_ptr + r).to(tl.int64)
+    _prefetch(pf_ptr, pf_bytes, PF, BULK)
+    _pdl(PDL)
+    HALF: tl.constexpr = HD // 2
+    G: tl.constexpr = NQ // NKV
+    d = tl.arange(0, HD)
+    is_k = h >= NQ
+    g = h - NQ
+    ph = tl.where(is_k, NXQ + g, h % NXQ)
+    col = ph * HD
+    base = qkv_ptr + r * qkv_row
+    keep = t + RING >= T   # only a call's last RING tokens go to the rings (no write races)
+    if CCGQA:
+        if PRE:
+            x = _conv_pre(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cb1, cacc, ca0, ca1, cw0, cw1, r, t, b, slot,
+                          pos, col, HD, KS, KC)
+        else:
+            x = _conv(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cs_ptr, csb_ptr, ch_ptr, chb_ptr, h, r, t, b, slot,
+                      pos, col, HD, KS, KC)
+        if QK_MEAN:
+            if is_k:
+                qm = tl.zeros([HD], dtype=tl.float32)
+                for jh in tl.static_range(G):
+                    qm += tl.load(base + ((g * G + jh) % NXQ) * HD + d).to(tl.float32)
+                x += 0.5 * (qm / G + tl.load(base + col + d).to(tl.float32))
+            else:
+                x += 0.5 * (tl.load(base + col + d).to(tl.float32)
+                            + tl.load(base + (NXQ + h // G) * HD + d).to(tl.float32))
+    else:
+        x = tl.load(base + col + d).to(tl.float32)
+    gain = tl.load(tl.where(is_k, kg_ptr + g * HD, qg_ptr) + d)
+    x = x * tl.rsqrt(tl.sum(x * x, 0) / HD + eps) * gain
+    if not CCGQA:   # plain QK-norm rounds to bf16 before RoPE
+        x = x.to(tl.bfloat16).to(tl.float32)
+    lo, hi = tl.split(tl.permute(tl.reshape(x, [2, HALF]), (1, 0)))
+    if ROPE:
+        i2 = tl.arange(0, HALF)
+        c = tl.load(cos_ptr + pos * HALF + i2)
+        s = tl.load(sin_ptr + pos * HALF + i2)
+        lo, hi = lo * c - hi * s, hi * c + lo * s
+    lo_i = tl.arange(0, HALF)
+    if is_k:
+        dst = kc_ptr + b * kc_b + g * kc_h + slot * HD
+    else:
+        dst = q_ptr + r * q_row + h * HD
+    tl.store(dst + lo_i, lo.to(q_ptr.dtype.element_ty))
+    tl.store(dst + HALF + lo_i, hi.to(q_ptr.dtype.element_ty))
+    if is_k:
+        ch = g * HD + d
+        own = tl.load(base + V_OFF + ch)
+        v = own
+        if VSHIFT:
+            VH: tl.constexpr = NKV * HD // 2
+            upper = ch >= VH
+            if t >= 1:
+                prev = tl.load(base - qkv_row + V_OFF + ch, mask=upper, other=0.0)
+            else:
+                prev = tl.load(vh_ptr + b * vh_b + ((slot - 1) % RING) * vh_l + ch - VH,
+                               mask=upper & (slot >= 1), other=0.0)
+            v = tl.where(upper, tl.where(pos > 0, prev, 0.0), own)
+            tl.store(vh_ptr + b * vh_b + (slot % RING) * vh_l + ch - VH, own, mask=upper & keep)
+        tl.store(vc_ptr + b * kc_b + g * kc_h + slot * HD + d, v.to(vc_ptr.dtype.element_ty))
+    if CCGQA:
+        if ((h < NXQ) | is_k) & keep:
+            tl.store(lat_ptr + b * lat_b + (slot % RING) * lat_l + col + d, tl.load(base + col + d))
+
+
+def qk_decode(qkv, lat, vh, pos, slot, cos, sin, conv, qg, kg, q, kc, vc, T: int, eps: float, *, nq, nkv,
+              nxq, hd, v_off, ks, kc_taps, ccgqa, qk_mean, vshift, rope, pf=None) -> None:
+    R = qkv.shape[0]
+    cs, csb, ch, chb = conv if conv is not None else (qkv,) * 4
+    _qk_decode_kernel[(R, nq + nkv)](
+        qkv, qkv.stride(0), lat if lat is not None else qkv, lat.stride(0) if lat is not None else 0,
+        lat.stride(1) if lat is not None else 0, vh if vh is not None else qkv,
+        vh.stride(0) if vh is not None else 0, vh.stride(1) if vh is not None else 0,
+        pos, slot, cos, sin, cs, csb, ch, chb, qg, kg, q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1),
+        T, eps, NQ=nq, NKV=nkv, NXQ=nxq, HD=hd, V_OFF=v_off, KS=ks, KC=kc_taps, CCGQA=ccgqa,
+        QK_MEAN=qk_mean, VSHIFT=vshift, ROPE=rope, num_warps=8, PDL=PDL, launch_pdl=PDL, **_pf(pf))
+
+
+@triton.jit
+def _tile_lat(qkv_ptr, qkv_row, r0, t, tm, pos, col, d, T, LAG: tl.constexpr):
+    """``[BT, HD]`` pre-conv latents ``LAG`` tokens back of a prompt tile."""
+    ok = tm & (pos >= LAG) & (t >= LAG)
+    return tl.load(qkv_ptr + (r0 + t - LAG)[:, None] * qkv_row + col + d[None, :], mask=ok[:, None],
+                   other=0.0).to(tl.float32)
+
+
+@triton.jit
+def _qk_prefill_kernel(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, vh_ptr, vh_b, vh_l, pos_ptr, cos_ptr, sin_ptr,
+                       cs_ptr, csb_ptr, ch_ptr, chb_ptr, qg_ptr, kg_ptr, q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h,
+                       T, eps, NQ: tl.constexpr, NKV: tl.constexpr, NXQ: tl.constexpr, HD: tl.constexpr,
+                       V_OFF: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr, QK_MEAN: tl.constexpr,
+                       VSHIFT: tl.constexpr, ROPE: tl.constexpr, BT: tl.constexpr):
+    """``qk_decode`` for a prompt, ``BT`` tokens of one virtual head per program (taps on tensor cores)."""
+    tt = tl.program_id(0)
+    h = tl.program_id(1)
+    b = tl.program_id(2).to(tl.int64)
+    HALF: tl.constexpr = HD // 2
+    G: tl.constexpr = NQ // NKV
+    t = tt * BT + tl.arange(0, BT)
+    tm = t < T
+    r0 = b * T
+    rows = r0 + t
+    pos = tl.load(pos_ptr + rows, mask=tm, other=0).to(tl.int64)
+    d = tl.arange(0, HD)
+    is_k = h >= NQ
+    g = h - NQ
+    ph = tl.where(is_k, NXQ + g, h % NXQ)
+    col = ph * HD
+    b1 = tl.load(csb_ptr + h * HD + d).to(tl.float32)
+    x = tl.zeros([BT, HD], dtype=tl.float32) + tl.load(chb_ptr + h * HD + d).to(tl.float32)[None, :]
+    for j in tl.static_range(KC):
+        u = tl.zeros([BT, HD], dtype=tl.float32) + b1[None, :]
+        for i in tl.static_range(KS):
+            a = tl.load(cs_ptr + (h * HD + d) * KS + i).to(tl.float32)
+            u += a[None, :] * _tile_lat(qkv_ptr, qkv_row, r0, t, tm, pos, col, d, T, i + j)
+        w2 = tl.load(ch_ptr + ((h * KC + j) * HD + d[:, None]) * HD + d[None, :]).to(tl.bfloat16)
+        u_hi = u.to(tl.bfloat16)
+        x = tl.dot(u_hi, w2, x)
+        x = tl.dot((u - u_hi.to(tl.float32)).to(tl.bfloat16), w2, x)
+    base = qkv_ptr + rows[:, None] * qkv_row
+    if QK_MEAN:
+        if is_k:
+            qm = tl.zeros([BT, HD], dtype=tl.float32)
+            for jh in tl.static_range(G):
+                qm += tl.load(base + ((g * G + jh) % NXQ) * HD + d[None, :], mask=tm[:, None], other=0.0).to(
+                    tl.float32)
+            x += 0.5 * (qm / G + tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0).to(tl.float32))
+        else:
+            x += 0.5 * (tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0).to(tl.float32)
+                        + tl.load(base + (NXQ + h // G) * HD + d[None, :], mask=tm[:, None], other=0.0).to(
+                            tl.float32))
+    gain = tl.load(tl.where(is_k, kg_ptr + g * HD, qg_ptr) + d)
+    x = x * tl.rsqrt(tl.sum(x * x, 1) / HD + eps)[:, None] * gain[None, :]
+    lo, hi = tl.split(tl.permute(tl.reshape(x, [BT, 2, HALF]), (0, 2, 1)))
+    i2 = tl.arange(0, HALF)
+    if ROPE:
+        c = tl.load(cos_ptr + pos[:, None] * HALF + i2[None, :])
+        sn = tl.load(sin_ptr + pos[:, None] * HALF + i2[None, :])
+        lo, hi = lo * c - hi * sn, hi * c + lo * sn
+    if is_k:
+        dst = kc_ptr + b * kc_b + g * kc_h + t[:, None] * HD
+    else:
+        dst = q_ptr + rows[:, None] * q_row + h * HD
+    tl.store(dst + i2[None, :], lo.to(q_ptr.dtype.element_ty), mask=tm[:, None])
+    tl.store(dst + HALF + i2[None, :], hi.to(q_ptr.dtype.element_ty), mask=tm[:, None])
+    keep = tm & (t + RING >= T)
+    if is_k:
+        ch = g * HD + d
+        own = tl.load(base + V_OFF + ch[None, :], mask=tm[:, None], other=0.0)
+        v = own
+        if VSHIFT:
+            VH: tl.constexpr = NKV * HD // 2
+            upper = (ch >= VH)[None, :]
+            prev = tl.load(base - qkv_row + V_OFF + ch[None, :], mask=tm[:, None] & upper & (t >= 1)[:, None],
+                           other=0.0)
+            v = tl.where(upper, tl.where((pos > 0)[:, None], prev, 0.0), own)
+            tl.store(vh_ptr + b * vh_b + (t % RING)[:, None] * vh_l + ch[None, :] - VH, own,
+                     mask=keep[:, None] & upper)
+        tl.store(vc_ptr + b * kc_b + g * kc_h + t[:, None] * HD + d[None, :], v.to(vc_ptr.dtype.element_ty),
+                 mask=tm[:, None])
+    if (h < NXQ) | is_k:
+        tl.store(lat_ptr + b * lat_b + (t % RING)[:, None] * lat_l + col + d[None, :],
+                 tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0), mask=keep[:, None])
+
+
+def qk_prefill(qkv, lat, vh, pos, cos, sin, conv, qg, kg, q, kc, vc, T: int, eps: float, *, nq, nkv, nxq, hd,
+               v_off, ks, kc_taps, qk_mean, vshift, rope) -> None:
+    """CCGQA ``qk_decode`` for ``B`` prompts of ``T`` tokens starting at slot 0."""
+    B = qkv.shape[0] // T
+    cs, csb, ch, chb = conv
+    bt = 32 if hd <= 128 else 16  # 64-token tiles spill
+    _qk_prefill_kernel[(triton.cdiv(T, bt), nq + nkv, B)](
+        qkv, qkv.stride(0), lat, lat.stride(0), lat.stride(1), vh if vh is not None else qkv,
+        vh.stride(0) if vh is not None else 0, vh.stride(1) if vh is not None else 0, pos, cos, sin,
+        cs, csb, ch, chb, qg, kg, q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1), T, eps, NQ=nq, NKV=nkv,
+        NXQ=nxq, HD=hd, V_OFF=v_off, KS=ks, KC=kc_taps, QK_MEAN=qk_mean, VSHIFT=vshift, ROPE=rope, BT=bt,
+        num_warps=4)
+
+
+# grid (span, KV head, sequence); the last span program merges the partials and finishes the rows
+
+
+@triton.jit
+def _spans(start, slot, NSPL: tl.constexpr, BL: tl.constexpr, MINB: tl.constexpr = 1):
+    """``(span, n_spans)`` for keys ``[start, slot]``, sized on the device."""
+    n = slot + 1 - start
+    span = tl.maximum(tl.cdiv(tl.cdiv(n, NSPL), BL), MINB) * BL
+    return span, tl.where(n > 0, tl.cdiv(n, span), 0)
+
+
+@triton.jit
+def _merge(ml_ptr, acc_ptr, base0, n_sp, hrow, HD: tl.constexpr, G_PAD: tl.constexpr, GH_PAD: tl.constexpr,
+           NSPL: tl.constexpr, CH: tl.constexpr):
+    """Merge the span partials of the query heads ``hrow`` -> bf16-rounded output."""
+    d = tl.arange(0, HD)
+    si = tl.arange(0, NSPL)
+    m = tl.load(ml_ptr + (base0 + si)[:, None] * 2 * G_PAD + hrow[None, :], mask=(si < n_sp)[:, None],
+                other=float("-inf"))
+    l = tl.load(ml_ptr + (base0 + si)[:, None] * 2 * G_PAD + G_PAD + hrow[None, :], mask=(si < n_sp)[:, None],
+                other=0.0)
+    mg = tl.max(m, 0)
+    mg = tl.where(mg == float("-inf"), 0.0, mg)
+    lsum = tl.sum(tl.exp2(m - mg[None, :]) * l, 0)
+    o = tl.zeros([GH_PAD, HD], dtype=tl.float32)
+    for c0 in range(0, n_sp, CH):
+        cs = c0 + tl.arange(0, CH)
+        cm = cs < n_sp
+        mc = tl.load(ml_ptr + (base0 + cs)[:, None] * 2 * G_PAD + hrow[None, :], mask=cm[:, None],
+                     other=float("-inf"))
+        a = tl.load(acc_ptr + ((base0 + cs)[:, None, None] * G_PAD + hrow[None, :, None]) * HD
+                    + d[None, None, :], mask=cm[:, None, None], other=0.0)
+        o += tl.sum(tl.exp2(mc - mg[None, :])[:, :, None] * a, 0)
+    o = o / tl.maximum(lsum, 1e-30)[:, None]
+    return o.to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _finish(ml_ptr, acc_ptr, r, g, b, n_sp, slot, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row,
+            lam_off, gate_off, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr,
+            HD: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr, GATE: tl.constexpr):
+    """Merge row ``r``'s spans for KV head ``g`` and finish DiffAttn V2 + XSA (+ gate) into ``y``."""
+    base0 = (r * NKV + g) * NSPL
+    d = tl.arange(0, HD)
+    if DIFF:
+        GH: tl.constexpr = G // 2
+        GH_PAD: tl.constexpr = G_PAD // 2
+        i = tl.arange(0, GH_PAD)
+        im = i < GH
+        y0 = _merge(ml_ptr, acc_ptr, base0, n_sp, 2 * i, HD, G_PAD, GH_PAD, NSPL, CH)
+        y1 = _merge(ml_ptr, acc_ptr, base0, n_sp, 2 * i + 1, HD, G_PAD, GH_PAD, NSPL, CH)
+        lam = tl.load(qkv_ptr + r * qkv_row + lam_off + g * GH + i, mask=im, other=0.0).to(tl.float32)
+        y = y0 - tl.sigmoid(lam)[:, None] * y1
+    else:
+        GH: tl.constexpr = G
+        GH_PAD: tl.constexpr = G_PAD
+        i = tl.arange(0, GH_PAD)
+        im = i < GH
+        y = _merge(ml_ptr, acc_ptr, base0, n_sp, i, HD, G_PAD, GH_PAD, NSPL, CH)
+    if XSA:
+        v = tl.load(vc_ptr + b * kc_b + g * kc_h + slot * HD + d).to(tl.float32)
+        vn = v / tl.maximum(tl.sqrt(tl.sum(v * v, 0)), 1e-12)
+        y = y - tl.sum(y * vn[None, :], 1)[:, None] * vn[None, :]
+    oh = g * GH + i
+    if GATE == 1:   # per channel
+        gt = tl.load(qkv_ptr + r * qkv_row + gate_off + oh[:, None] * HD + d[None, :], mask=im[:, None],
+                     other=0.0).to(tl.float32)
+        y = y * tl.sigmoid(gt)
+    elif GATE == 2:   # per head
+        gt = tl.load(qkv_ptr + r * qkv_row + gate_off + oh, mask=im, other=0.0).to(tl.float32)
+        y = y * tl.sigmoid(gt)[:, None]
+    tl.store(y_ptr + r * y_row + oh[:, None] * HD + d[None, :], y.to(y_ptr.dtype.element_ty), mask=im[:, None])
+
+
+@triton.jit
+def _fence():
+    """``fence.acq_rel.gpu`` in every thread (CUDA's ``__threadfence``)."""
+    tl.inline_asm_elementwise("fence.acq_rel.gpu; mov.u32 $0, 0;", "=r", [], dtype=tl.int32, is_pure=False,
+                              pack=1)
+
+
+@triton.jit
+def _seq_keys(start_ptr, slot_ptr, r0, T, win, WIN: tl.constexpr):
+    """A sequence's key range ``[start, last]``."""
+    start = tl.load(start_ptr + r0)
+    if WIN:
+        start = tl.maximum(start, tl.load(slot_ptr + r0) - win + 1)
+    return start, tl.load(slot_ptr + r0 + T - 1)
+
+
+@triton.jit
+def _unit(x, tol):
+    """``x / |x|``; 0 where ``|x| <= tol`` (the port's ``_unit``)."""
+    n = tl.sqrt(tl.sum(x * x, 0))
+    return tl.where(n > tol, x / tl.maximum(n, tol), 0.0)
+
+
+@triton.jit
+def _iha_finish(ml_ptr, acc_ptr, rt, g, b, n_sp, vs0, r_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
+                y_row, gate_off, P: tl.constexpr, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr,
+                G_PAD: tl.constexpr, HD: tl.constexpr, CH: tl.constexpr, XSA: tl.constexpr, GATE: tl.constexpr):
+    """IHA finish: merge, collapse with ``R``, Gram-Schmidt XSA and the gate for one token row."""
+    d = tl.arange(0, HD)
+    i = tl.arange(0, G_PAD)
+    im = i < G
+    oh = g * G + i
+    y = tl.zeros([G_PAD, HD], dtype=tl.float32)
+    for p in tl.static_range(P):
+        yp = _merge(ml_ptr, acc_ptr, ((rt * P + p) * NKV + g) * NSPL, n_sp, i, HD, G_PAD, G_PAD, NSPL, CH)
+        rr = tl.load(r_ptr + oh * P + p, mask=im, other=0.0).to(tl.float32)
+        y += rr[:, None] * yp
+    y = y.to(tl.bfloat16).to(tl.float32)
+    if XSA:
+        vb = vc_ptr + b * kc_b + g * kc_h + vs0 * HD
+        e0 = tl.zeros([HD], dtype=tl.float32)
+        e1 = tl.zeros([HD], dtype=tl.float32)
+        e2 = tl.zeros([HD], dtype=tl.float32)
+        for p in tl.static_range(P):
+            w = _unit(tl.load(vb + p * HD + d).to(tl.float32), 1e-12)
+            w = w - tl.sum(w * e0, 0) * e0
+            w = w - tl.sum(w * e1, 0) * e1
+            w = w - tl.sum(w * e2, 0) * e2
+            e = _unit(w, 1e-2)
+            y = y - tl.sum(y * e[None, :], 1)[:, None] * e[None, :]
+            if p == 0:
+                e0 = e
+            elif p == 1:
+                e1 = e
+            elif p == 2:
+                e2 = e
+    if GATE == 1:
+        gt = tl.load(qkv_ptr + rt * qkv_row + gate_off + oh[:, None] * HD + d[None, :], mask=im[:, None],
+                     other=0.0).to(tl.float32)
+        y = y * tl.sigmoid(gt)
+    elif GATE == 2:
+        gt = tl.load(qkv_ptr + rt * qkv_row + gate_off + oh, mask=im, other=0.0).to(tl.float32)
+        y = y * tl.sigmoid(gt)[:, None]
+    tl.store(y_ptr + rt * y_row + oh[:, None] * HD + d[None, :], y.to(y_ptr.dtype.element_ty), mask=im[:, None])
+
+
+@triton.jit
+def _attn_kernel(q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h, slot_ptr, start_ptr, ml_ptr, acc_ptr, cnt_ptr,
+                 qkv_ptr, qkv_row, y_ptr, y_row, T, qk_scale, lam_off, gate_off, win, r_ptr,
+                 NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr, QR: tl.constexpr,
+                 HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr,
+                 GATE: tl.constexpr, FINISH: tl.constexpr, WIN: tl.constexpr = False, IHA: tl.constexpr = 0,
+                 MINB: tl.constexpr = 1, PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
+    """Flash-decoding attention; ``FINISH``: the last span program finishes the row."""
+    sp = tl.program_id(0)
+    g = tl.program_id(1)
+    b = tl.program_id(2).to(tl.int64)
+    r0 = b * T
+    start, last = _seq_keys(start_ptr, slot_ptr, r0, T, win, WIN)
+    span, n_sp = _spans(start, last, NSPL, BL, MINB)
+    lo = start + sp * span
+    hi = tl.minimum(last + 1, lo + span)
+    qi = tl.arange(0, QR)
+    t = qi // G
+    j = qi % G
+    qm = t < T
+    lim = tl.load(slot_ptr + r0 + t, mask=qm, other=-1)
+    _prefetch(pf_ptr, pf_bytes, PF, BULK)
+    _pdl(PDL)
+    if lo < hi:
+        d = tl.arange(0, HD)
+        q = tl.load(q_ptr + (r0 + t)[:, None] * q_row + (g * G + j)[:, None] * HD + d[None, :],
+                    mask=qm[:, None], other=0.0)
+        m_i = tl.full([QR], float("-inf"), tl.float32)
+        l_i = tl.zeros([QR], dtype=tl.float32)
+        acc = tl.zeros([QR, HD], dtype=tl.float32)
+        kbase = b * kc_b + g * kc_h
+        for k0 in range(lo, hi, BL):
+            kk = k0 + tl.arange(0, BL)
+            km = kk < hi
+            k = tl.load(kc_ptr + kbase + kk[:, None] * HD + d[None, :], mask=km[:, None], other=0.0)
+            s = tl.dot(q, tl.trans(k)) * qk_scale                       # base-2 logits
+            ok = km[None, :] & (kk[None, :] <= lim[:, None])
+            if WIN:
+                ok = ok & (kk[None, :] > lim[:, None] - win)
+            s = tl.where(ok, s, float("-inf"))
+            m_new = tl.maximum(m_i, tl.max(s, 1))
+            m_ok = tl.where(m_new == float("-inf"), 0.0, m_new)         # a row with no key yet
+            alpha = tl.exp2(m_i - m_ok)
+            p = tl.exp2(s - m_ok[:, None])
+            l_i = l_i * alpha + tl.sum(p, 1)
+            v = tl.load(vc_ptr + kbase + kk[:, None] * HD + d[None, :], mask=km[:, None], other=0.0)
+            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
+            m_i = m_new
+        base = ((r0 + t) * NKV + g) * NSPL + sp
+        tl.store(ml_ptr + base * 2 * G_PAD + j, m_i, mask=qm)
+        tl.store(ml_ptr + base * 2 * G_PAD + G_PAD + j, l_i, mask=qm)
+        tl.store(acc_ptr + (base * G_PAD + j)[:, None] * HD + d[None, :], acc, mask=qm[:, None])
+        if FINISH:
+            _fence()
+            tl.debug_barrier()
+            old = tl.atomic_add(cnt_ptr + b * NKV + g, 1, sem="acq_rel", scope="gpu")
+            if old == n_sp - 1:
+                _fence()
+                tl.atomic_xchg(cnt_ptr + b * NKV + g, 0, sem="relaxed", scope="gpu")
+                if IHA:
+                    _iha_finish(ml_ptr, acc_ptr, b, g, b, n_sp, tl.load(slot_ptr + r0), r_ptr, qkv_ptr, qkv_row,
+                                vc_ptr, kc_b, kc_h, y_ptr, y_row, gate_off, IHA, NKV, NSPL, G, G_PAD, HD, CH,
+                                XSA, GATE)
+                else:
+                    _finish(ml_ptr, acc_ptr, r0, g, b, n_sp, last, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
+                            y_row, lam_off, gate_off, NKV, NSPL, G, G_PAD, HD, CH, DIFF, XSA, GATE)
+
+
+@triton.jit
+def _finish_kernel(ml_ptr, slot_ptr, start_ptr, acc_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row, T,
+                   lam_off, gate_off, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr,
+                   HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr,
+                   GATE: tl.constexpr, MINB: tl.constexpr = 1, PDL: tl.constexpr = False):
+    """The finish of a multi-row (verify / short prompt) attention, a program per (row, KV head)."""
+    _pdl(PDL)
+    r = tl.program_id(0).to(tl.int64)
+    g = tl.program_id(1)
+    b = r // T
+    _, n_sp = _spans(tl.load(start_ptr + b * T), tl.load(slot_ptr + b * T + T - 1), NSPL, BL, MINB)
+    _finish(ml_ptr, acc_ptr, r, g, b, n_sp, tl.load(slot_ptr + r), qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
+            y_row, lam_off, gate_off, NKV, NSPL, G, G_PAD, HD, CH, DIFF, XSA, GATE)
+
+
+@triton.jit
+def _iha_finish_kernel(ml_ptr, slot_ptr, start_ptr, acc_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row,
+                       Tt, win, r_ptr, gate_off, P: tl.constexpr, NKV: tl.constexpr, NSPL: tl.constexpr,
+                       G: tl.constexpr, G_PAD: tl.constexpr, HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr,
+                       XSA: tl.constexpr, GATE: tl.constexpr, WIN: tl.constexpr, MINB: tl.constexpr = 1,
+                       PDL: tl.constexpr = False):
+    """IHA's finish of a multi-token call, a program per (token row, KV head)."""
+    _pdl(PDL)
+    rt = tl.program_id(0).to(tl.int64)
+    g = tl.program_id(1)
+    b = rt // Tt
+    start, last = _seq_keys(start_ptr, slot_ptr, b * Tt * P, Tt * P, win, WIN)
+    _, n_sp = _spans(start, last, NSPL, BL, MINB)
+    _iha_finish(ml_ptr, acc_ptr, rt, g, b, n_sp, tl.load(slot_ptr + rt * P), r_ptr, qkv_ptr, qkv_row, vc_ptr,
+                kc_b, kc_h, y_ptr, y_row, gate_off, P, NKV, NSPL, G, G_PAD, HD, CH, XSA, GATE)
+
+
+# tests set 32 / False so a 1-row step and a verify match bitwise
+QR_MIN = 16
+FUSED_FINISH = True
+# smallest span, in 64-key blocks
+ATTN_MINB = 1
+
+
+def attention(q, kc, vc, slot, start, qkv, y, ml, acc, cnt, T: int, *, nq, nkv, hd, nspl, lam_off,
+              gate_off, diff, xsa, gate, win: int = 0, iha_p: int = 0, iha_r=None, pf=None, bl: int = 64,
+              warps: int | None = None) -> None:
+    """Flash-decoding attention for ``R = B * T`` rows (``win``: sliding window; ``iha_p``: pseudo-rows)."""
+    if iha_p:
+        assert not diff, "IHA with DiffAttn V2: no fused finish"
+        T = T * iha_p
+    R = q.shape[0]
+    G = nq // nkv
+    g_pad = max(16, triton.next_power_of_2(G))
+    qr = max(QR_MIN, triton.next_power_of_2(T * G))
+    ch = 8 if g_pad <= 16 else 4
+    _attn_kernel[(nspl, nkv, R // T)](q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1), slot, start, ml, acc,
+                                      cnt, qkv, qkv.stride(0), y, y.stride(0), T, hd ** -0.5 * LOG2E, lam_off,
+                                      gate_off, int(win), iha_r if iha_p else q,
+                                      NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad, QR=qr, HD=hd, BL=bl,
+                                      CH=ch, DIFF=diff, XSA=xsa, GATE=gate,
+                                      FINISH=T == max(iha_p, 1) and FUSED_FINISH, WIN=win > 0, IHA=iha_p,
+                                      MINB=ATTN_MINB,
+                                      num_warps=warps or (4 if qr <= 32 else 8), num_stages=2, PDL=PDL, launch_pdl=PDL,
+                                      **_pf(pf))
+    if iha_p:
+        if T > iha_p or not FUSED_FINISH:
+            _iha_finish_kernel[(R // iha_p, nkv)](ml, slot, start, acc, qkv, qkv.stride(0), vc, kc.stride(0),
+                                                  kc.stride(1), y, y.stride(0), T // iha_p, int(win), iha_r,
+                                                  gate_off, P=iha_p, NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad,
+                                                  HD=hd, BL=bl, CH=ch, XSA=xsa, GATE=gate, WIN=win > 0, MINB=ATTN_MINB,
+                                                  num_warps=4, PDL=PDL, launch_pdl=PDL)
+        return
+    if T > 1 or not FUSED_FINISH:
+        _finish_kernel[(R, nkv)](ml, slot, start, acc, qkv, qkv.stride(0), vc, kc.stride(0), kc.stride(1), y,
+                                 y.stride(0), T, lam_off, gate_off, NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad, HD=hd,
+                                 BL=bl, CH=ch, DIFF=diff, XSA=xsa, GATE=gate, MINB=ATTN_MINB, num_warps=4, PDL=PDL,
+                                 launch_pdl=PDL)
+
+
+@triton.jit
+def _iha_mix_kernel(q_ptr, q_row, kl_ptr, vl_ptr, kl_b, kl_h, aq_ptr, ak_ptr, av_ptr, qg_ptr, kg_ptr, temp_ptr,
+                    pos_ptr, slot_ptr, cos_ptr, sin_ptr, qv_ptr, qv_row, kc_ptr, vc_ptr, kc_b, kc_h, T, eps,
+                    P: tl.constexpr, NQ: tl.constexpr, NQ_PAD: tl.constexpr, NKV: tl.constexpr, HD: tl.constexpr,
+                    ROPE: tl.constexpr, PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
+    """One (row, pseudo-head p): IHA mixes, norm, temperature and RoPE, then cache / ``qv`` writes."""
+    r = tl.program_id(0).to(tl.int64)
+    p = tl.program_id(1)
+    b = r // T
+    HALF: tl.constexpr = HD // 2
+    rslot = tl.load(slot_ptr + r).to(tl.int64)
+    pos = tl.load(pos_ptr + r).to(tl.int64) * P + p
+    slot = rslot * P + p
+    d = tl.arange(0, HD)
+    i2 = tl.arange(0, HALF)
+    h = tl.arange(0, NQ_PAD)
+    hm = h < NQ
+    at = tl.load(aq_ptr + (h[None, :] * NQ + h[:, None]) * P + p, mask=hm[:, None] & hm[None, :], other=0.0)
+    qgain = tl.load(qg_ptr + d)
+    kgain = tl.load(kg_ptr + d)
+    if ROPE:
+        c = tl.load(cos_ptr + pos * HALF + i2)
+        s = tl.load(sin_ptr + pos * HALF + i2)
+    _prefetch(pf_ptr, pf_bytes, PF, BULK)
+    _pdl(PDL)
+    # queries: one tensor-core tile
+    q = tl.load(q_ptr + r * q_row + h[:, None] * HD + d[None, :], mask=hm[:, None], other=0.0)
+    x = tl.dot(at.to(tl.bfloat16), q.to(tl.bfloat16)).to(tl.bfloat16).to(tl.float32)   # at[h, m]
+    x = x * tl.rsqrt(tl.sum(x * x, 1) / HD + eps)[:, None] * qgain[None, :]
+    x = x.to(tl.bfloat16).to(tl.float32)
+    lo, hi = tl.split(tl.permute(tl.reshape(x, [NQ_PAD, 2, HALF]), (0, 2, 1)))
+    if ROPE:
+        lo, hi = lo * c[None, :] - hi * s[None, :], hi * c[None, :] + lo * s[None, :]
+    dst = qv_ptr + (r * P + p) * qv_row + h[:, None] * HD + i2[None, :]
+    tl.store(dst, lo.to(qv_ptr.dtype.element_ty), mask=hm[:, None])
+    tl.store(dst + HALF, hi.to(qv_ptr.dtype.element_ty), mask=hm[:, None])
+    # keys / values: CUDA-core sums
+    src = b * kl_b + rslot * HD
+    for g in tl.static_range(NKV):
+        k = tl.zeros([HD], dtype=tl.float32)
+        v = tl.zeros([HD], dtype=tl.float32)
+        for m in tl.static_range(NKV):
+            ak = tl.load(ak_ptr + (m * NKV + g) * P + p).to(tl.float32)
+            av = tl.load(av_ptr + (m * NKV + g) * P + p).to(tl.float32)
+            k += ak * tl.load(kl_ptr + src + m * kl_h + d).to(tl.float32)
+            v += av * tl.load(vl_ptr + src + m * kl_h + d).to(tl.float32)
+        k = k.to(tl.bfloat16).to(tl.float32)
+        k = (k * tl.rsqrt(tl.sum(k * k, 0) / HD + eps) * kgain).to(tl.bfloat16).to(tl.float32)
+        k = (k * tl.load(temp_ptr + g)).to(tl.bfloat16).to(tl.float32)
+        klo, khi = tl.split(tl.permute(tl.reshape(k, [2, HALF]), (1, 0)))
+        if ROPE:
+            klo, khi = klo * c - khi * s, khi * c + klo * s
+        kd = kc_ptr + b * kc_b + g * kc_h + slot * HD
+        tl.store(kd + i2, klo.to(kc_ptr.dtype.element_ty))
+        tl.store(kd + HALF + i2, khi.to(kc_ptr.dtype.element_ty))
+        tl.store(vc_ptr + b * kc_b + g * kc_h + slot * HD + d, v.to(vc_ptr.dtype.element_ty))
+
+
+def iha_mix(q, kl, vl, aq, ak, av, qg, kg, temp, pos, slot, cos, sin, qv, kc, vc, T: int, eps: float, *,
+            P: int, nq: int, nkv: int, hd: int, rope: bool, pf=None) -> None:
+    """``q`` ``(R, nq, hd)`` latent queries, ``kl`` / ``vl`` ``(B, nkv, L, hd)`` latent keys / values at the
+    rows' slots, mixes ``(H, H, P)`` bf16, ``qv`` ``(R * P, nq, hd)``, caches ``(B, nkv, L * P, hd)``."""
+    R = q.shape[0]
+    _iha_mix_kernel[(R, P)](q, q.stride(0), kl, vl, kl.stride(0), kl.stride(1), aq, ak, av, qg, kg, temp, pos,
+                            slot, cos, sin, qv, qv.stride(0), kc, vc, kc.stride(0), kc.stride(1), T, eps,
+                            P=P, NQ=nq, NQ_PAD=max(16, triton.next_power_of_2(nq)), NKV=nkv, HD=hd, ROPE=rope,
+                            num_warps=4, PDL=PDL, launch_pdl=PDL, **_pf(pf))
+
+
+@triton.jit
+def _head_gate_kernel(y_ptr, y_row, b_ptr, h_ptr, h_row, D, RK: tl.constexpr, BLOCK: tl.constexpr,
+                      PDL: tl.constexpr = False):
+    """Gated byte head input ``[h, sigmoid(u . h + b) * P h]``."""
+    _pdl(PDL)
+    r = tl.program_id(0).to(tl.int64)
+    j = tl.arange(0, BLOCK)
+    jm = j < RK
+    a = tl.load(y_ptr + r * y_row + RK).to(tl.float32) + tl.load(b_ptr).to(tl.float32)
+    gt = tl.sigmoid(a.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32)
+    yp = tl.load(y_ptr + r * y_row + j, mask=jm, other=0.0).to(tl.float32)
+    tl.store(h_ptr + r * h_row + D + j, (gt * yp).to(h_ptr.dtype.element_ty), mask=jm)
+
+
+def head_gate(y, bias, h, D: int, rank: int) -> None:
+    _head_gate_kernel[(y.shape[0],)](y, y.stride(0), bias, h, h.stride(0), D, RK=rank,
+                                     BLOCK=triton.next_power_of_2(rank), num_warps=4, PDL=PDL, launch_pdl=PDL)
+
+
+@triton.jit
+def _sample_kernel(x_ptr, x_row, V, inv_temp, top_k, top_p, min_p, seed_ptr, excl_ptr, y_ptr, pd_ptr, u_ptr,
+                   pres_ptr, pres_row, ids_ptr, penalty, T,
+                   TOP_K: tl.constexpr, TOP_P: tl.constexpr, MIN_P: tl.constexpr, EXCL: tl.constexpr,
+                   REP: tl.constexpr, GREEDY: tl.constexpr, T_MAX: tl.constexpr, BLOCK: tl.constexpr,
+                   ITERS: tl.constexpr):
+    """One row: repetition penalty, temperature, top-k, top-p, min-p (thresholds by bisection), then a
+    draw."""
+    r = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, BLOCK)
+    msk = offs < V
+    x = tl.load(x_ptr + r * x_row + offs, mask=msk, other=float("-inf"))
+    if REP:
+        b = r // T
+        t = r % T
+        seen = tl.load(pres_ptr + b * pres_row + offs, mask=msk, other=0) != 0
+        for i in tl.static_range(1, T_MAX):
+            if i <= t:
+                seen = seen | (offs == tl.load(ids_ptr + b * T + i))
+        x = tl.where(seen, tl.where(x < 0, x * penalty, x / penalty), x)
+    if GREEDY:
+        tl.store(y_ptr + r, tl.argmax(x, 0).to(tl.int64))
+    else:
+        x = x * inv_temp
+        if TOP_K:
+            lo = tl.min(tl.where(msk, x, float("inf")), 0)
+            hi = tl.max(x, 0)
+            for _ in range(ITERS):   # count(x >= lo) >= k > count(x >= hi)
+                mid = lo + (hi - lo) * 0.5
+                ok = tl.sum((x >= mid).to(tl.int32), 0) >= top_k
+                lo = tl.where(ok, mid, lo)
+                hi = tl.where(ok, hi, mid)
+            kth = tl.min(tl.where(x >= lo, x, float("inf")), 0)
+            x = tl.where(x >= kth, x, float("-inf"))
+        m = tl.max(x, 0)
+        p = tl.exp(x - m)                      # max 1
+        z = tl.sum(p, 0)
+        if TOP_P:
+            lo = tl.zeros([], dtype=tl.float32)
+            hi = tl.full([], 1.0, tl.float32)  # mass(p >= lo) >= top_p * z > mass(p >= hi) (or hi = 1)
+            goal = top_p * z
+            for _ in range(ITERS):
+                mid = lo + (hi - lo) * 0.5
+                ok = tl.sum(tl.where(p >= mid, p, 0.0), 0) >= goal
+                lo = tl.where(ok, mid, lo)
+                hi = tl.where(ok, hi, mid)
+            tau = tl.min(tl.where(p >= lo, p, float("inf")), 0)
+            p = tl.where(p >= tau, p, 0.0)
+        if MIN_P:
+            p = tl.where(p >= min_p, p, 0.0)   # max is 1: min_p * max = min_p
+        seed = tl.load(seed_ptr)
+        if EXCL:
+            e = tl.load(excl_ptr + r)
+            pd = tl.sum(tl.where(offs == e, p, 0.0), 0) / tl.sum(p, 0)
+            tl.store(pd_ptr + r, pd)
+            tl.store(u_ptr + r, tl.rand(seed + 7919, r.to(tl.int32)))
+            p = tl.where(offs == e, 0.0, p)
+        cdf = tl.cumsum(p, 0)
+        u = tl.rand(seed, r.to(tl.int32)) * tl.max(cdf, 0)
+        tok = tl.min(tl.where((cdf > u) & (p > 0), offs, BLOCK), 0)
+        tl.store(y_ptr + r, tl.where(tok < V, tok, tl.argmax(p, 0)).to(tl.int64))
+
+
+def sample(x: torch.Tensor, seed: torch.Tensor, y: torch.Tensor, *, temperature: float, top_k: int, top_p: float,
+           min_p: float, excl: torch.Tensor | None = None, pd: torch.Tensor | None = None,
+           u: torch.Tensor | None = None, greedy: bool = False, penalty: float = 1.0,
+           pres: torch.Tensor | None = None, ids: torch.Tensor | None = None, T: int = 1) -> None:
+    """Draw ``y`` ``(R,)`` from fp32 logits ``x`` ``(R, V)``."""
+    R, V = x.shape
+    block = triton.next_power_of_2(V)
+    rep = penalty != 1.0
+    _sample_kernel[(R,)](x, x.stride(0), V, 1.0 / temperature, int(top_k), float(top_p), float(min_p), seed,
+                         excl if excl is not None else y, y, pd if pd is not None else x,
+                         u if u is not None else x, pres.view(torch.uint8) if rep else y,
+                         pres.stride(0) if rep else 0, ids if rep else y, float(penalty), T,
+                         TOP_K=bool(top_k) and top_k < V, TOP_P=top_p < 1.0, MIN_P=bool(min_p),
+                         EXCL=excl is not None, REP=rep, GREEDY=greedy, T_MAX=T, BLOCK=block, ITERS=40,
+                         num_warps=32 if block >= 16384 else 8)
+
+
+kernels = sys.modules[__name__]   # ``from .luau_kernels import kernels as K``: transformers'
+                                  # remote code only follows ``from .module import ...`` lines

modeling_luau_coder.py

@@ -1,3 +1,9 @@
+"""LuauCoder for transformers (``trust_remote_code=True``).
+
+Same math and parameter names as ``lm/model.py``. ``generate()`` uses the fused
+engine in ``luau_infer.py`` when its settings allow.
+"""
+
 from __future__ import annotations
 
 import copy
@@ -35,16 +41,11 @@ __all__ = ["LuauCoderCache", "LuauCoderForCausalLM", "LuauCoderModel", "LuauCode
 
 _V5 = int(transformers.__version__.split(".")[0]) >= 5
 IGNORE_INDEX = -100
-_HW_TANH = torch.version.hip is None  # PTX tanh.approx in the fused CE; ROCm falls back to sigmoid
-
-
-# fused elementwise
+_HW_TANH = torch.version.hip is None  # PTX tanh.approx; ROCm uses sigmoid
 
 
 class _Fused:
-    """``fn`` through ``torch.compile`` on CUDA (one generated kernel per
-    direction); eager on CPU, inside an outer compile, when disabled, or after
-    a compilation failure (warned once)."""
+    """``fn`` through ``torch.compile`` on CUDA; eager otherwise or after a compile failure."""
 
     def __init__(self, fn):
         self.fn = fn
@@ -60,7 +61,7 @@ class _Fused:
             return self.compiled(*args)
         except torch.OutOfMemoryError:
             raise
-        except Exception as exc:  # no working Inductor / Triton toolchain here
+        except Exception as exc:
             self.failed = True
             warnings.warn(f"torch.compile of {self.fn.__name__} failed ({exc!r}); running it eagerly",
                           stacklevel=2)
@@ -71,9 +72,7 @@ _CONSTS: dict = {}
 
 
 def _cached(tensors: tuple, fn):
-    """``fn()``, a value derived from the parameters ``tensors``. Without grad
-    (inference, outside ``torch.compile``) it is cached until one of them changes:
-    decoding would otherwise rebuild these fp32 weights at every token."""
+    """``fn()`` derived from ``tensors``; cached without grad until one of them changes."""
     if torch.is_grad_enabled() or torch.compiler.is_compiling():
         return fn()
     key = tuple((t.data_ptr(), t._version) for t in tensors)
@@ -117,10 +116,7 @@ def _qk_rope_fn(q: Tensor, k: Tensor, qw: Tensor | None, kw: Tensor | None, cos:
 
 def _diff_xsa_fn(y: Tensor, lam: Tensor | None, v: Tensor, nh: int, nkv: int, xsa: bool,
                  gate: Tensor | None = None) -> Tensor:
-    """``y`` ``(B, T, nq, d)`` attention output, ``lam`` ``(B, T, nh)`` lambda
-    logits (None: no DiffAttn), ``v`` ``(B, T, nkv, d)`` the tokens' own values,
-    ``gate`` output-gate logits ``(B, T, nh * d)`` per channel or ``(B, T, nh)`` per head
-    (None: no gate), applied after XSA."""
+    """DiffAttn -> XSA -> output gate (each optional)."""
     B, T, _, hd = y.shape
     if lam is not None:
         y = y.view(B, T, nh, 2, hd).float()
@@ -144,8 +140,7 @@ def _unit(x: Tensor, tol: float) -> Tensor:
 
 def _iha_diff_xsa(y: Tensor, lam: Tensor | None, vp: Tensor, nh: int, nkv: int, xsa: bool,
                   gate: Tensor | None = None) -> Tensor:
-    """``_diff_xsa_fn`` for IHA's collapsed heads (training: ``_iha_diff_xsa_gate``): XSA removes the
-    span of the token's own pseudo-values ``vp`` ``(B, T, P, nkv, d)``, by modified Gram-Schmidt."""
+    """``_diff_xsa_fn`` for IHA's collapsed heads: XSA removes the span of the pseudo-values."""
     B, T, _, hd = y.shape
     if lam is not None:
         y = y.view(B, T, nh, 2, hd).float()
@@ -168,8 +163,7 @@ def _iha_diff_xsa(y: Tensor, lam: Tensor | None, vp: Tensor, nh: int, nkv: int,
 
 
 def _iha_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scaling: float) -> Tensor:
-    """GQA-native attention with an explicit mask: ``q`` ``(B, nq, Tq, d)``, ``k`` / ``v`` ``(B, nkv, Tk, d)``,
-    ``mask`` ``(B|1, 1, 1, Tq, Tk)`` bool -> ``(B, Tq, nq, d)``; fp32 scores and softmax."""
+    """GQA attention with an explicit bool mask; fp32 scores and softmax."""
     B, nq, Tq, d = q.shape
     nkv = k.shape[1]
     s = torch.matmul(q.reshape(B, nkv, nq // nkv, Tq, d).float(), k.float().unsqueeze(2).transpose(-1, -2)) * scaling
@@ -178,10 +172,9 @@ def _iha_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scaling: float
 
 
 def _iha_mix_init(alpha: Tensor, stride: int) -> None:
-    """Training's init: pseudo-head ``p`` of head ``h`` starts as base head ``(h + p * stride) % H``
-    plus a small Gaussian (``alpha`` ``(H, H, P)``; ``stride`` 2 for DiffAttn's query heads)."""
+    """Training's IHA mix init."""
     H, _, P = alpha.shape
-    if getattr(alpha, "_is_hf_initialized", False):   # loaded from a checkpoint
+    if getattr(alpha, "_is_hf_initialized", False):
         return
     with torch.no_grad():
         alpha.normal_(std=0.1 / H ** 0.5)
@@ -198,20 +191,15 @@ def _swiglu_fn(gu: Tensor) -> Tensor:
 def _ccgqa_fn(qk: Tensor, tail: Tensor, v: Tensor, before: Tensor | None, pos: Tensor, W: Tensor,
               b: Tensor, q_gain: Tensor, k_gain: Tensor, cos: Tensor | None, sin: Tensor | None, nq: int,
               nkv: int, eps: float, qk_mean: bool) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor | None]:
-    """CCGQA ``q`` ``(B, T, nq, d)``, ``k`` / ``v`` ``(B, T, nkv, d)`` from the q|k latents
-    ``qk`` ``(B, T, nxq+nkv, d)`` (``nxq < nq``: VHA's physical query heads) and the raw
-    values, in fp32 like ``ccgqa_qk_reference``. ``tail`` holds the ``K-1`` latents
-    before ``qk`` and ``before`` the previous token's value-shift half (None: no value
-    shift); both come back updated. ``W`` / ``b`` / gains: ``_ccgqa_consts``. Tap ``m``
-    reads token ``t-m`` only if ``pos[t] >= m``."""
+    """CCGQA q / k / v from the q|k latents in fp32; ``tail`` / ``before`` carry decode state."""
     B, T, _, hd = qk.shape
     K, nxq = tail.shape[1] + 1, qk.shape[2] - nkv
     ext = torch.cat((tail, qk), 1)                            # (B, K-1+T, nxq+nkv, d)
 
-    def virtual(x: Tensor) -> Tensor:   # VHA: query head h reads physical head h % nxq
+    def virtual(x: Tensor) -> Tensor:
         return x if nxq == nq else torch.cat((x[:, :, :nxq].repeat(1, 1, nq // nxq, 1), x[:, :, nxq:]), 2)
 
-    win = virtual(ext).unfold(1, K, 1)                        # (B, T, H, d, K): slot j = lag K-1-j
+    win = virtual(ext).unfold(1, K, 1)  # slot j = lag K-1-j
     lag = torch.arange(K - 1, -1, -1, device=qk.device)
     win = torch.where((pos[..., None] >= lag)[:, :, None, None, :], win.float(), 0.0)
     H = win.shape[2]
@@ -229,7 +217,7 @@ def _ccgqa_fn(qk: Tensor, tail: Tensor, v: Tensor, before: Tensor | None, pos: T
     if before is None:
         return q, k, v.contiguous(), tail, None
     flat = v.reshape(B, T, nkv * hd)
-    half = nkv * hd // 2          # the W_V-bar half of the channels is the previous token's
+    half = nkv * hd // 2
     shifted = torch.where((pos > 0)[..., None], torch.cat((before, flat[:, :, half:]), 1)[:, :T], 0)
     v = torch.cat((flat[:, :, :half], shifted), -1).view(B, T, nkv, hd)
     return q, k, v, tail, flat[:, -1:, half:]
@@ -271,7 +259,7 @@ _ccgqa = _Fused(_ccgqa_fn)
 _MAX_HIDDEN = 8192     # whole row in registers
 _MAX_SOURCES = 32
 _SPAN_BWD = dict(R=64, num_warps=2)
-_TPP = 8               # tokens per program in the per-token backward
+_TPP = 8
 
 if HAVE_TRITON:
 
@@ -280,8 +268,7 @@ if HAVE_TRITON:
                        s_st, s_sk, x_st, o_st, st_st, eps,
                        D: tl.constexpr, NB: tl.constexpr, NB_PAD: tl.constexpr, BLOCK: tl.constexpr,
                        ADD: tl.constexpr, PART: tl.constexpr):
-        """One token: online softmax over the NB sources + x, one read of each
-        (``ADD``: plus their sum, the residual stream; ``PART`` off: x isn't routed)."""
+        """One token: online softmax over the NB sources + x."""
         t = tl.program_id(0).to(tl.int64)
         offs = tl.arange(0, BLOCK)
         msk = offs < D
@@ -327,9 +314,7 @@ if HAVE_TRITON:
                              D: tl.constexpr, NB: tl.constexpr, NB_PAD: tl.constexpr,
                              HAS_RES: tl.constexpr, TPP: tl.constexpr, BLOCK: tl.constexpr,
                              ADD: tl.constexpr):
-        """TPP tokens: <g, v_i> for every entry (one read each), the full
-        gradient of x (plus the residual-add path's, ``HAS_RES``), the scalars
-        the deferred source gradient needs, and the x part of dw."""
+        """Per-token mix backward."""
         pid = tl.program_id(0)
         offs = tl.arange(0, BLOCK)
         msk = offs < D
@@ -375,8 +360,7 @@ if HAVE_TRITON:
                             s_st, s_sk, ds_st, ds_sk, N, coef_c, coef_q,
                             D: tl.constexpr, NB: tl.constexpr, NC: tl.constexpr, NC_PAD: tl.constexpr,
                             R: tl.constexpr, CD: tl.constexpr):
-        """dS for every source of one block from all NC mixes that read it (one
-        read + one write per source), plus the S part of each mix's dw."""
+        """dS for every source of one block from all mixes that read it."""
         pid_t = tl.program_id(0)
         cols = tl.program_id(1) * CD + tl.arange(0, CD)
         mc = cols < D
@@ -428,9 +412,7 @@ def _ar_forward(x: Tensor, S: Tensor, w: Tensor, eps: float, additive: bool,
 
 def attn_res_reference(x: Tensor, sources: Tensor, norm_weight: Tensor, query: Tensor,
                        eps: float, additive: bool = False, partial: bool = True) -> Tensor:
-    """Eager fp32 mix; ``x`` ``(N, D)``, ``sources`` ``(N, k, D)``. ``additive``: Delta
-    Block, the mix added to the stream ``x + sum(sources)``. ``partial=False``: ``x``
-    takes no part in the softmax."""
+    """Eager fp32 mix reference."""
     values = torch.cat((sources, x.unsqueeze(1)), dim=1).float()
     keys = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + eps)
     logits = (keys * (norm_weight.float() * query.float())).sum(-1)
@@ -447,19 +429,18 @@ class _SpanState:
 
     def __init__(self, n_calls: int, additive: bool = False, partial: bool = True):
         self.n_calls = n_calls
-        self.additive, self.partial = additive, partial   # Delta Block; x is routed
-        self.w: Tensor | None = None          # (n_calls, D) fp32 score weights
+        self.additive, self.partial = additive, partial
+        self.w: Tensor | None = None
         self.gbuf: Tensor | None = None       # (n_calls, N, D) output grads, filled in backward
         self.coef: Tensor | None = None       # (n_calls, 3, N, k) fp32 a / b / e
         self.dw_x: Tensor | None = None       # (n_calls, D) fp32 x part of dw
-        self.res: dict[int, Tensor] = {}      # call -> residual-add gradient of its x
-        self.mixed: set[int] = set()          # calls whose mix backward has run
+        self.res: dict[int, Tensor] = {}
+        self.mixed: set[int] = set()
         self.done = 0
 
 
 class _SpanSources(torch.autograd.Function):
-    """Identity on the sources ``S``; its backward (after every mix of the
-    block) builds ``dS`` and the score-weight gradients in one pass."""
+    """Identity on the sources; its backward builds ``dS`` once per block."""
 
     @staticmethod
     def forward(ctx, S: Tensor, state: _SpanState, *weights: Tensor) -> Tensor:
@@ -499,9 +480,7 @@ class _SpanSources(torch.autograd.Function):
 
 
 class _ResidualAdd(torch.autograd.Function):
-    """``x + y`` where ``x`` is also read by mix ``call``: the backward hands
-    the gradient of ``x`` to that mix's backward kernel, which sums it into
-    ``dx``, instead of autograd summing the two in a separate pass."""
+    """``x + y``; the gradient of ``x`` is summed inside mix ``call``'s backward kernel."""
 
     @staticmethod
     def forward(ctx, x: Tensor, y: Tensor, state: _SpanState, call: int) -> Tensor:
@@ -511,7 +490,7 @@ class _ResidualAdd(torch.autograd.Function):
     @staticmethod
     def backward(ctx, g: Tensor):
         st = ctx.state
-        if ctx.call in st.mixed:  # the mix already ran its backward: let autograd sum
+        if ctx.call in st.mixed:
             return g, g, None, None
         st.res[ctx.call] = g
         return None, g, None, None
@@ -558,14 +537,7 @@ class _DepthMix(torch.autograd.Function):
 
 
 class AttnResSpan:
-    """The sources ``S`` ``(N, k, D)`` of one block and the mixes reading them.
-
-    ``sites[c]`` is the ``(norm, query)`` pair of the ``c``-th mix that reads
-    ``S``. With grad enabled every site must be called exactly once per
-    forward. Paths: Triton with autograd (training), Triton forward only (no
-    grad: inference / generation), eager fp32 reference (CPU, no Triton,
-    ``fused_kernels=False``). ``additive``: Delta Block mixes; ``partial``: ``x`` is routed too.
-    """
+    """The sources of one block and the mixes reading them; each site is called once per forward."""
 
     def __init__(self, S: Tensor, sites: list[tuple[Tensor, Tensor]], eps: float, fused: bool,
                  additive: bool = False, partial: bool = True):
@@ -606,8 +578,6 @@ class AttnResSpan:
         return x + y
 
 
-# fused LM head + CE
-
 def softcap_logits(logits: Tensor, cap: float | None) -> Tensor:
     """``cap * tanh(logits / cap)`` (``final_logit_softcapping``); identity for None / 0."""
     return cap * torch.tanh(logits / cap) if cap else logits
@@ -617,7 +587,7 @@ if HAVE_TRITON:
 
     @triton.jit
     def _tanh(x, HW: tl.constexpr):
-        if HW:  # one MUFU op (sm_75+); rel. error 2^-11, under the bf16 rounding of the logits
+        if HW:
             return tl.inline_asm_elementwise("tanh.approx.f32 $0, $1;", "=r,r", [x],
                                              dtype=tl.float32, is_pure=True, pack=1)
         return 2.0 * tl.sigmoid(2.0 * x) - 1.0
@@ -626,9 +596,7 @@ if HAVE_TRITON:
     def _ce_fwd_bwd_kernel(logits_ptr, stride_row, labels_ptr, loss_ptr, scale_ptr, V, cap, inv_cap,
                            IGNORE: tl.constexpr, SOFTCAP: tl.constexpr, HW_TANH: tl.constexpr,
                            BLOCK: tl.constexpr):
-        """NLL of one row, then its logits overwritten by ``(softmax - onehot) * scale``
-        (with ``SOFTCAP``: of ``cap * tanh(logits / cap)``, times ``1 - tanh^2``).
-        Base 2 throughout so ``log2(e)`` folds into the ``ex2`` argument."""
+        """NLL of one row; its logits are overwritten with their gradient (base 2)."""
         LOG2E: tl.constexpr = 1.4426950408889634
         LN2: tl.constexpr = 0.6931471805599453
         row = tl.program_id(0).to(tl.int64)
@@ -637,8 +605,7 @@ if HAVE_TRITON:
         valid = label != IGNORE
         cap2 = cap * LOG2E
         safe_label = tl.where(valid, label, 0)
-        # the target stays out of the sum-exp so 1 - p_t is s_nt / s, not
-        # 1 - round(p_t) (Softmax Collapse, Barron et al. 2026)
+        # target kept out of the sum-exp (Softmax Collapse)
         m = -float("inf")
         s_nt = 0.0
         for off in range(0, V, BLOCK):
@@ -659,12 +626,12 @@ if HAVE_TRITON:
         else:
             target = target * LOG2E
         s = tl.exp2(target - m) + s_nt
-        u = 1.0 + s_nt  # target is the max: loss = log1p(s_nt) (Goldberg's log1p)
+        u = 1.0 + s_nt
         log1p = tl.where(u == 1.0, s_nt, tl.log(u) * (s_nt / (u - 1.0)))
         loss = tl.where(target >= m, log1p, (m + tl.log2(s) - target) * LN2)
         tl.store(loss_ptr + row, tl.where(valid, loss, 0.0))
         r = tl.where(valid, tl.load(scale_ptr), 0.0) / s
-        g_t = -s_nt * r  # -(1 - p_t) * scale: the row sums to zero
+        g_t = -s_nt * r
         for off in range(0, V, BLOCK):
             cols = off + tl.arange(0, BLOCK)
             msk = cols < V
@@ -689,9 +656,7 @@ def _accumulate_dw(dw: Tensor, dlogits: Tensor, h: Tensor) -> None:
 
 
 class _LinearCrossEntropy(torch.autograd.Function):
-    """Chunked ``CE(h @ W.T, y)`` (sum over valid targets times ``scale``):
-    per chunk one logits GEMM, one Triton pass that writes ``dlogits`` over
-    the logits, and the two gradient GEMMs, so the backward only rescales."""
+    """Chunked ``CE(h @ W.T, y)`` that computes gradients during the forward."""
 
     @staticmethod
     def forward(ctx, hidden: Tensor, weight: Tensor, labels: Tensor, scale: Tensor, chunk: int,
@@ -748,14 +713,8 @@ def fused_linear_cross_entropy(hidden: Tensor, weight: Tensor, labels: Tensor,
                                      chunk, softcap)
 
 
-# CCGQA
-
-
 def compose_conv(w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> tuple[Tensor, Tensor]:
-    """Depthwise taps ``w1`` ``(H*hd, ks)`` + ``b1`` and per-head taps ``w2``
-    ``(H, kc, hd_in, hd_out)`` + ``b2`` ``(H, hd)`` (taps in lag order, 0 = current
-    token) -> one masked causal conv: fp32 ``W`` ``(H, ks+kc-1, hd, hd)``, ``b``
-    ``(H, hd)``. Exact inside a document (``lm/kernels/ccgqa.py``)."""
+    """Compose the depthwise and per-head convs into one masked causal conv."""
     H, kc, hd, _ = w2.shape
     ks = w1.shape[1]
     a, w2f = w1.float().view(H, hd, ks, 1), w2.float()
@@ -766,10 +725,7 @@ def compose_conv(w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> tuple[Tensor
 
 
 class LuauCoderCache(DynamicCache):
-    """``DynamicCache`` plus CCGQA's per-layer decode state: the last ``K-1``
-    pre-conv q|k latents and the previous token's raw value-shift half.
-    Rewinding (``crop``: assisted / speculative decoding) is not supported
-    once that state exists."""
+    """``DynamicCache`` plus CCGQA's conv / value-shift decode state (no ``crop`` once it exists)."""
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -816,18 +772,7 @@ class _InPlaceStates(dict):
 
 
 class LuauCoderStaticCache:
-    """Preallocated cache for CUDA-graph decoding (``generate.py``): K/V buffers
-    ``(B, nkv, L, d)`` per layer, CCGQA's conv / value-shift state, and which
-    buffer slots hold real keys. Every tensor keeps its address and the write
-    slots live on the device, so a captured decode step replays as is.
-
-    Before each forward the caller sets ``idx`` (the slots of this call's
-    tokens, ``(T,)``) and marks them in ``key_valid``, and passes
-    ``position_ids`` (each token's position in its document; a left-padded
-    row starts at its first real token). Query ``t`` attends to the valid
-    slots ``<= idx[t]``. Only ``generate.py``'s loop drives it: not for
-    ``generate()``.
-    """
+    """Preallocated cache for ``generate.py``'s CUDA-graph loop; the caller sets ``idx`` / ``key_valid``."""
 
     is_static = True
 
@@ -871,12 +816,8 @@ def _conv_states(cache) -> dict:
     return states
 
 
-# RoPE
-
-
 def _rope_settings(config: LuauCoderConfig) -> tuple[float, dict | None]:
-    """``(theta, yarn params or None)`` from a v4 (``rope_theta`` + ``rope_scaling``)
-    or v5 (``rope_parameters``) config."""
+    """``(theta, yarn params or None)`` from a v4 or v5 config."""
     params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
     theta = float(params.get("rope_theta", getattr(config, "rope_theta", 500000.0)))
     rtype = params.get("rope_type", params.get("type"))
@@ -886,8 +827,7 @@ def _rope_settings(config: LuauCoderConfig) -> tuple[float, dict | None]:
 
 
 def rope_inv_freq(head_dim: int, theta: float, yarn: dict | None) -> tuple[Tensor, float]:
-    """Inverse frequencies ``(head_dim/2,)`` fp64 and the cos/sin scale; YaRN
-    "NTK-by-parts" as in the training stack and transformers' ``rope_type: yarn``."""
+    """Inverse frequencies fp64 and the cos/sin scale (YaRN as in training)."""
     dim = head_dim
     pos_freqs = theta ** (torch.arange(0, dim, 2, dtype=torch.float64, device="cpu") / dim)
     extrap = 1.0 / pos_freqs
@@ -915,16 +855,13 @@ def rope_inv_freq(head_dim: int, theta: float, yarn: dict | None) -> tuple[Tenso
 
 
 def lerope_inv_freq(head_dim: int, theta: float, alpha: list[float], yarn: dict | None) -> tuple[Tensor, float]:
-    """LeRoPE frequencies fp64 and the cos/sin scale, as the training stack's ``lerope_freqs``: the
-    RoPE (+ YaRN) frequencies times ``exp(alpha)`` (transformers' Mistral Large 4 convention)."""
+    """LeRoPE frequencies and the cos/sin scale."""
     inv, mscale = rope_inv_freq(head_dim, theta, yarn)
     return inv * torch.tensor(alpha, dtype=torch.float64, device="cpu").exp(), mscale
 
 
 class LuauCoderRotaryEmbedding(nn.Module):
-    """fp32 cos/sin tables (computed in fp64 like training), gathered by
-    ``position_ids``. The tables are plain attributes, not buffers, so a
-    ``model.to(bfloat16)`` never rounds them."""
+    """fp32 cos/sin tables gathered by ``position_ids`` (plain attributes so ``.to(bf16)`` skips them)."""
 
     def __init__(self, config: LuauCoderConfig):
         super().__init__()
@@ -946,13 +883,8 @@ class LuauCoderRotaryEmbedding(nn.Module):
         return self._cos[position_ids], self._sin[position_ids]
 
 
-# modules
-
-
 class LuauCoderRMSNorm(nn.Module):
-    """``learnable=False`` (q/k norm): the gain is a persistent buffer, fixed at 1 as in training
-    (same key and shape, so checkpoints and the fused kernels are unchanged; not a parameter, so it
-    is neither counted nor trained)."""
+    """``learnable=False``: the gain is a buffer fixed at 1."""
 
     def __init__(self, dim: int, eps: float, fuse: bool = True, learnable: bool = True):
         super().__init__()
@@ -993,10 +925,7 @@ def eager_attention_forward(module: nn.Module, query: Tensor, key: Tensor, value
 
 
 def _decode_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scaling: float) -> Tensor:
-    """One query token against a static cache, GQA-native: ``q`` ``(B, nq, 1, d)``,
-    ``k`` / ``v`` ``(B, nkv, L, d)``, ``mask`` ``(B, 1, 1, L)`` bool -> ``(B, 1, nq, d)``.
-    Each KV head is read once (the backends' masked path repeats K/V per query head,
-    which dominates a decode step at a long cache); fp32 scores and softmax."""
+    """Single-token decode attention against a static cache, reading each KV head once."""
     B, nq, _, d = q.shape
     nkv = k.shape[1]
     s = torch.matmul(q.reshape(B, nkv, nq // nkv, d).float(), k.float().transpose(-1, -2)) * scaling
@@ -1005,15 +934,13 @@ def _decode_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scaling: fl
 
 
 class LuauCoderAttention(nn.Module):
-    """DiffAttn V2 + XSA. Query heads ``2i`` / ``2i+1`` are adjacent, so both
-    fall in the same GQA group and share K and V."""
+    """DiffAttn V2 + XSA attention."""
 
     def __init__(self, config: LuauCoderConfig, layer_idx: int):
         super().__init__()
         self.config, self.layer_idx = config, layer_idx
         self.nh, self.nkv, self.hd = config.num_attention_heads, config.num_key_value_heads, config.head_dim
         self.nq = config.num_query_heads
-        # VHA: qkv makes nxq physical q heads; query head h reads physical h % nxq
         self.nxq = getattr(config, "num_q_latent_heads", self.nq)
         self.diff, self.xsa = config.diff_v2, config.xsa
         self.num_key_value_groups = self.nq // self.nkv
@@ -1024,20 +951,15 @@ class LuauCoderAttention(nn.Module):
         self.rope = layer_idx not in (getattr(config, "nope_layers", None) or ())
         self.iha = config.iha_layer(layer_idx)
         D = config.hidden_size
-        # output gate: per channel, its W_g rows follow q|k|v in the fused qkv weight;
-        # per head (LFM3), a separate dim -> nh gate_proj
         self.g_rows = self.nh * self.hd if config.output_gate == "elementwise" else 0
         self.qkv = nn.Linear(D, (self.nxq + 2 * self.nkv) * self.hd + self.g_rows, bias=False)
         self.lambda_proj = nn.Linear(D, self.nh, bias=False) if self.diff else None
         self.gate_proj = nn.Linear(D, self.nh, bias=False) if config.output_gate == "headwise" else None
         self.o_proj = nn.Linear(self.nh * self.hd, D, bias=False)
-        # gains fixed at 1 like training (lm/model.py: RMSNorm(learnable=False))
         self.q_norm = LuauCoderRMSNorm(self.hd, self.eps, learnable=False) if config.qk_norm else None
         self.k_norm = LuauCoderRMSNorm(self.hd, self.eps, learnable=False) if config.qk_norm else None
         self.ccgqa = getattr(config, "ccgqa", False)
         if self.ccgqa:
-            # q|k latents (q heads, then k heads): depthwise causal conv taps
-            # conv_seq[c, i] and per-head dense taps conv_head[h, j, c_in, c_out], lag order
             H, ks, kc = self.nq + self.nkv, config.ccgqa_seq_kernel, config.ccgqa_head_kernel
             self.conv_taps = ks + kc - 1
             self.qk_mean, self.value_shift = config.ccgqa_qk_mean, config.ccgqa_value_shift
@@ -1045,7 +967,7 @@ class LuauCoderAttention(nn.Module):
             self.conv_seq_bias = nn.Parameter(torch.zeros(H * self.hd))
             self.conv_head = nn.Parameter(torch.empty(H, kc, self.hd, self.hd))
             self.conv_head_bias = nn.Parameter(torch.zeros(H, self.hd))
-            self.k_temp = nn.Parameter(torch.zeros(self.nkv))  # key scale exp(t), per KV head
+            self.k_temp = nn.Parameter(torch.zeros(self.nkv))  # key scale exp(t)
         if self.iha:
             P = self.iha_p = config.iha_pseudo_heads
             self.iha_q = nn.Parameter(torch.empty(self.nq, self.nq, P))
@@ -1063,7 +985,7 @@ class LuauCoderAttention(nn.Module):
         conv_pos = kwargs.pop("conv_pos", None)
         iha_ctx = kwargs.pop("iha_ctx", None)
         rope = self.rope
-        if not self.rope or self.iha:   # IHA layers rotate later, at the virtual positions
+        if not self.rope or self.iha:
             cos = sin = None
         if self.ccgqa:
             qk, v = qkv.split([self.nxq + self.nkv, self.nkv], dim=2)
@@ -1078,7 +1000,6 @@ class LuauCoderAttention(nn.Module):
             y, vp = self._iha_attend(q, k, v, iha_ctx, rope, past_key_values)
             probs = None
         else:
-            # (B, H, T, d) views: flash transposes them straight back, sdpa takes the strides
             q, k_all, v_all = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
             if past_key_values is not None:
                 k_all, v_all = past_key_values.update(k_all, v_all, self.layer_idx)
@@ -1089,7 +1010,7 @@ class LuauCoderAttention(nn.Module):
                 fn = eager_attention_forward if impl in (None, "eager") else ALL_ATTENTION_FUNCTIONS[impl]
                 y, probs = fn(self, q, k_all, v_all, attention_mask, dropout=0.0, scaling=self.scaling,
                               position_ids=position_ids, **kwargs)       # (B, T, nq, d)
-        if self.diff and self.gate_proj is not None:   # one GEMM for both nh-wide projections
+        if self.diff and self.gate_proj is not None:
             lam, gate = F.linear(x, torch.cat((self.lambda_proj.weight, self.gate_proj.weight))).split(self.nh, -1)
         else:
             lam = self.lambda_proj(x) if self.diff else None
@@ -1104,11 +1025,7 @@ class LuauCoderAttention(nn.Module):
 
     def _iha_attend(self, q: Tensor, k: Tensor, v: Tensor, ctx: dict, rope: bool,
                     cache: Cache | None) -> tuple[Tensor, Tensor]:
-        """Training's ``Attention._iha_attend`` on latent, unrotated q / k / v: pseudo-heads, QK-norm,
-        temperature, RoPE at virtual positions (``ctx``: ``LuauCoderModel._iha_context``), then
-        attention over the interleaved tokens, with the pseudo-keys in the KV cache as extra heads
-        ``(B, nkv * P, T, d)``. Returns the collapsed ``(B, T, nq, d)`` and the pseudo-values
-        ``(B, T, P, nkv, d)``."""
+        """IHA attention with the pseudo-keys cached as extra heads."""
         B, T, _, hd = q.shape
         P, nq, nkv = self.iha_p, self.nq, self.nkv
 
@@ -1129,7 +1046,7 @@ class LuauCoderAttention(nn.Module):
         if cache is not None:
             keys, vals = cache.update(keys, vals, self.layer_idx)
 
-        def interleave(t: Tensor) -> Tensor:   # (B, nkv * P, T', d) -> (B, nkv, Tk * P, d), the tokens in reach
+        def interleave(t: Tensor) -> Tensor:
             t = t[:, :, ctx["first"]:]
             return t.reshape(B, nkv, P, -1, hd).transpose(2, 3).reshape(B, nkv, -1, hd)
 
@@ -1142,8 +1059,7 @@ class LuauCoderAttention(nn.Module):
 
     def _ccgqa_qkv(self, qk: Tensor, v: Tensor, cos: Tensor | None, sin: Tensor | None,
                    pos: Tensor | None, cache: Cache | None) -> tuple[Tensor, Tensor, Tensor]:
-        """CCGQA ``q`` / ``k`` / ``v`` (``_ccgqa_fn``), carrying the conv and value-shift
-        state in ``cache``. ``pos`` ``(B, T)``: each token's position in its document."""
+        """CCGQA q / k / v with conv and value-shift state in ``cache``."""
         B, T = qk.shape[:2]
         if pos is None:
             pos = torch.arange(T, device=qk.device).expand(B, T)
@@ -1161,9 +1077,7 @@ class LuauCoderAttention(nn.Module):
         return q, k, v
 
     def _ccgqa_consts(self) -> tuple[Tensor, Tensor, Tensor, Tensor]:
-        """fp32 composed conv ``W`` as ``(H, d*K, d)`` bmm operands (row ``c*K + j`` =
-        input channel ``c`` at window slot ``j``, lag ``K-1-j``), its bias, the q gain
-        and the k gain times ``exp(temp)``; cached without grad (``_cached``)."""
+        """fp32 composed conv weights, bias and gains (cached without grad)."""
         params = (self.conv_seq, self.conv_seq_bias, self.conv_head, self.conv_head_bias,
                   self.q_norm.weight, self.k_norm.weight, self.k_temp)
 
@@ -1171,7 +1085,6 @@ class LuauCoderAttention(nn.Module):
             W, b = compose_conv(*params[:4])
             H, K, hd, _ = W.shape
             W = W.flip(1).permute(0, 2, 1, 3).reshape(H, hd * K, hd)
-            # IHA applies the temperature after its mix and norm
             kg = params[5].float() if self.iha else params[5].float() * params[6].float().exp()[:, None]
             return W, b, params[4].float(), kg
 
@@ -1201,8 +1114,6 @@ class LuauCoderDecoderLayer(GradientCheckpointingLayer):
         self.ffn = LuauCoderMLP(config)
         self.attn_res = config.attn_res
         if self.attn_res:
-            # one pseudo-query and key-norm gain per sublayer input; layer 0's
-            # attention reads the embedding alone, so it has nothing to mix
             self.block_start = layer_idx % config.attn_res_block_size == 0
             self.attn_res_norm = nn.Parameter(torch.ones(D)) if layer_idx else None
             self.attn_res_query = nn.Parameter(torch.zeros(D)) if layer_idx else None
@@ -1223,10 +1134,7 @@ class LuauCoderDecoderLayer(GradientCheckpointingLayer):
 
     def _forward_attn_res(self, partial: Tensor, attn_span: AttnResSpan | None, mlp_span: AttnResSpan,
                           cos, sin, attention_mask, position_ids, past_key_values, **kwargs):
-        """``partial`` ``(B, T, D)``: the running sum of the current block's
-        sublayer outputs (at a block start: the finished block's, which
-        ``mlp_span`` already holds as its newest source). Returns the new
-        partial and this layer's attention input."""
+        """Block AttnRes layer; returns the new partial and this layer's attention input."""
         B, T, D = partial.shape
         N = B * T
         flat = partial.reshape(N, D)
@@ -1234,17 +1142,12 @@ class LuauCoderDecoderLayer(GradientCheckpointingLayer):
         a, _ = self.attn(self.attn_norm(h.view(B, T, D)), cos, sin, attention_mask, position_ids,
                          past_key_values, **kwargs)
         a = a.reshape(N, D).to(flat.dtype)
-        # each partial is read by exactly one mix and one residual add: the add's
-        # gradient is summed inside that mix's backward kernel
         flat = a if self.block_start else attn_span.residual(flat, a, self.attn_call)
         m = mlp_span.mix(flat, self.mlp_call)
         f = self.ffn(self.ffn_norm(m.view(B, T, D))).reshape(N, D).to(flat.dtype)
         return mlp_span.residual(flat, f, self.mlp_call).view(B, T, D), h.view(B, T, D)
 
 
-# models
-
-
 class LuauCoderPreTrainedModel(PreTrainedModel):
     config_class = LuauCoderConfig
     config: LuauCoderConfig
@@ -1262,9 +1165,7 @@ class LuauCoderPreTrainedModel(PreTrainedModel):
 
     @torch.no_grad()
     def _init_weights(self, module: nn.Module) -> None:
-        """The training stack's init: N(0, std) inputs, N(0, std / sqrt(2L))
-        residual-branch outputs, lambda logits at half scale, unit norms,
-        zero AttnRes pseudo-queries (uniform depth mix)."""
+        """The training stack's init."""
         cfg = self.config
         std = cfg.initializer_range
         out_std = std / math.sqrt(2 * cfg.num_hidden_layers)
@@ -1275,7 +1176,7 @@ class LuauCoderPreTrainedModel(PreTrainedModel):
                 _init.normal_(module.lambda_proj.weight, mean=0.0, std=0.5 * std)
             if module.gate_proj is not None:
                 _init.normal_(module.gate_proj.weight, mean=0.0, std=std)
-            if module.ccgqa:  # ZAYA init: conv weights N(0, std), zero biases, unit key scale
+            if module.ccgqa:  # ZAYA init
                 _init.normal_(module.conv_seq, mean=0.0, std=std)
                 _init.normal_(module.conv_head, mean=0.0, std=std)
                 for p in (module.conv_seq_bias, module.conv_head_bias, module.k_temp):
@@ -1295,6 +1196,9 @@ class LuauCoderPreTrainedModel(PreTrainedModel):
         elif isinstance(module, LuauCoderForCausalLM):
             if not cfg.tie_word_embeddings:
                 _init.normal_(module.lm_head.weight, mean=0.0, std=std)
+        elif isinstance(module, LuauCoderHeadGate):
+            _init.normal_(module.proj, mean=0.0, std=std)
+            _init.zeros_(module.bias)
         for name, p in module.named_parameters(recurse=False):
             if name.endswith("_res_norm"):
                 _init.ones_(p)
@@ -1303,13 +1207,10 @@ class LuauCoderPreTrainedModel(PreTrainedModel):
 
     @classmethod
     def _supports_default_dynamic_cache(cls) -> bool:
-        # generate() then leaves the cache to forward(), which makes a LuauCoderCache
-        # (so CCGQA's conv state follows beam reordering / batch expansion)
         return False
 
     def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
-        # AttnRes sources reach the layers as keyword arguments: reentrant
-        # checkpointing would silently drop their gradients
+        # reentrant checkpointing would drop the AttnRes sources' gradients
         gc_kwargs = dict(gradient_checkpointing_kwargs or {})
         gc_kwargs["use_reentrant"] = False
         return super().gradient_checkpointing_enable(gradient_checkpointing_kwargs=gc_kwargs, **kwargs)
@@ -1344,7 +1245,7 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
         self.norm = LuauCoderRMSNorm(config.hidden_size, config.rms_norm_eps, config.compile_elementwise)
         self.rotary_emb = LuauCoderRotaryEmbedding(config)
         if config.attn_res:
-            if not config.attn_res_delta_block:  # final mix over every block before the output norm
+            if not config.attn_res_delta_block:
                 self.out_res_norm = nn.Parameter(torch.ones(config.hidden_size))
                 self.out_res_query = nn.Parameter(torch.zeros(config.hidden_size))
             self._plan_attn_res()
@@ -1352,15 +1253,7 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
         self.post_init()
 
     def _plan_attn_res(self) -> None:
-        """Which block's sources each mix reads, and its call index there.
-
-        Block ``m`` (layers ``mS .. mS+S-1``) is read by its layers' MLP mixes
-        (call ``2j``), their attention mixes (``2j-1``, ``j >= 1``) and the
-        next block's first attention mix, which runs before that block is
-        pushed (``2S-1``); the last block's final slot is the output mix, which
-        Delta Block doesn't have. Sites are stored as ``(layer, kind)`` and resolved to parameters at
-        every forward (loading can replace the parameter objects).
-        """
+        """Assign each mix its source span and call index (resolved to parameters every forward)."""
         S, L = self.config.attn_res_block_size, self.config.num_hidden_layers
         n_spans = (L + S - 1) // S
         self._span_sites: list[list] = [[None] * (2 * min(S, L - m * S)) for m in range(n_spans)]
@@ -1394,10 +1287,7 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
 
     def _iha_context(self, position_ids: Tensor, pos: Tensor, attention_mask: Tensor | None,
                      past_len: int) -> dict:
-        """What every IHA layer of this forward shares: cos / sin ``(B|1, T * P, d/2)`` at virtual positions
-        ``position * P + p``, the first real token still in reach and the mask ``(B|1, 1, 1, T * P, Tk * P)``
-        over row indices: causal in virtual order, the window, padding and packed documents (``pos``
-        restarting). A row with no valid key (left padding) attends to itself: its softmax stays finite."""
+        """Shared IHA state: virtual-position cos / sin, window start and mask."""
         cfg = self.config
         P, W = cfg.iha_pseudo_heads, cfg.iha_window
         if attention_mask is not None and attention_mask.dim() != 2:
@@ -1436,8 +1326,8 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
         kwargs.pop("output_attentions", None)
         output_hidden_states = (output_hidden_states if output_hidden_states is not None
                                 else self.config.output_hidden_states)
-        if use_cache is None:  # no cache while training: it only costs memory and disables
-            use_cache = self.config.use_cache and not self.training  # packed-sequence masks
+        if use_cache is None:
+            use_cache = self.config.use_cache and not self.training
         if self.gradient_checkpointing and self.training:
             use_cache = False
         if inputs_embeds is None:
@@ -1448,7 +1338,6 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
             past_key_values = None
         B, T, D = inputs_embeds.shape
         if getattr(past_key_values, "is_static", False):
-            # no host syncs (CUDA-graph capture): the cache holds the slots and valid keys
             if position_ids is None or attention_mask is not None:
                 raise ValueError("a LuauCoderStaticCache takes position_ids and no attention_mask")
             if self.config.iha:
@@ -1468,8 +1357,6 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
         layer_kw = dict(cos=cos, sin=sin, attention_mask=mask, position_ids=position_ids,
                         past_key_values=past_key_values, **kwargs)
         if self.config.ccgqa or self.config.iha:
-            # document positions for the conv mask: position_ids restart per packed document,
-            # and a 2-D padding mask starts each row at its first real token
             pos = position_ids.long().expand(B, T)
             if attention_mask is not None and attention_mask.dim() == 2:
                 pos = torch.minimum(pos, (attention_mask.long().cumsum(-1) - 1)[:, -T:])
@@ -1480,7 +1367,7 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
 
         if self.config.attn_res:
             x = inputs_embeds
-            if torch.is_autocast_enabled(x.device.type):  # bf16 depth stream, like the sublayers
+            if torch.is_autocast_enabled(x.device.type):
                 x = x.to(torch.get_autocast_dtype(x.device.type))
             partial = x
             sources = x.reshape(B * T, 1, D)[:, :0]
@@ -1489,7 +1376,7 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
             delta = self.config.attn_res_delta_block
             route_x = self.config.attn_res_partial_source or not delta
             for i, blk in enumerate(self.blocks):
-                if blk.block_start:  # the finished block (or the embedding) becomes a source
+                if blk.block_start:
                     sources = torch.cat((sources, partial.reshape(B * T, 1, D)), dim=1)
                     sites = [self._site(*s) for s in self._span_sites[len(spans)]]
                     spans.append(AttnResSpan(sources, sites, eps, fused, delta, route_x))
@@ -1512,9 +1399,7 @@ class LuauCoderModel(LuauCoderPreTrainedModel):
                                        hidden_states=tuple(hidden) if output_hidden_states else None)
 
 
-# fused generate()
-
-# generation-config fields the fused decoder doesn't implement, with their neutral values
+# unsupported generation-config fields and their neutral values
 _FUSED_NEUTRAL = {
     "num_beams": (None, 1), "num_beam_groups": (None, 1), "diversity_penalty": (None, 0.0), "penalty_alpha": (None,),
     "dola_layers": (None,), "constraints": (None,), "force_words_ids": (None,), "prompt_lookup_num_tokens": (None,),
@@ -1528,8 +1413,6 @@ _FUSED_NEUTRAL = {
     "output_hidden_states": (None, False), "return_dict_in_generate": (None, False), "token_healing": (None, False),
     "max_time": (None,), "use_mtp": (None, False),
 }
-# generate() keyword arguments (model kwargs) the fused path accepts; anything else falls back.
-# (``use_cache`` / ``generation_mode`` are accepted and ignored: the fused path always caches.)
 _FUSED_MODEL_KWARGS = {"input_ids", "attention_mask", "use_cache", "tokenizer", "generation_mode"}
 
 
@@ -1538,13 +1421,13 @@ def _fused_fallback_reason(model, inputs, gc, model_kwargs, logits_processor, st
     """None if ``generate()`` can run on the fused decoder with these settings, else why not."""
     if not (HAVE_TRITON and model.device.type == "cuda") or model.training:
         return "needs CUDA + triton, eval mode"
-    if model.dtype != torch.bfloat16:   # its GEMMs round like bf16 ones: an fp32 model keeps its own precision
+    if model.dtype != torch.bfloat16:
         return f"a {model.dtype} model (the fused decoder runs bf16 models)"
     used = [k for k, v in extra.items() if v is not None and v is not False]
     if used:
         return f"unsupported generate() arguments {used}"
-    if logits_processor or stopping_criteria:
-        return "custom logits processors / stopping criteria"
+    if logits_processor:
+        return "custom logits processors"
     if set(model_kwargs) - _FUSED_MODEL_KWARGS:
         return f"model kwargs {sorted(set(model_kwargs) - _FUSED_MODEL_KWARGS)}"
     for name, ok in _FUSED_NEUTRAL.items():
@@ -1553,8 +1436,8 @@ def _fused_fallback_reason(model, inputs, gc, model_kwargs, logits_processor, st
     if not isinstance(inputs, Tensor) or inputs.dim() != 2 or inputs.is_floating_point():
         return "needs input_ids (B, T)"
     eos = gc.eos_token_id
-    if isinstance(eos, (list, tuple)) and len(set(eos)) > 1:
-        return "several EOS tokens"
+    if isinstance(eos, (list, tuple)) and len(set(eos)) > 8:
+        return "more than 8 EOS tokens"
     if gc.stop_strings and model_kwargs.get("tokenizer") is None:
         return "stop_strings without a tokenizer"
     mask = model_kwargs.get("attention_mask")
@@ -1565,17 +1448,8 @@ def _fused_fallback_reason(model, inputs, gc, model_kwargs, logits_processor, st
 
 
 class _FusedGenerationMixin:
-    """``generate()`` through the fused decoding engine (``luau_infer``: ~9 Triton kernels per
-    layer, each step one CUDA graph, prompt-lookup speculation with the same output) whenever
-    every setting is one it reproduces exactly: greedy or sampling (temperature, top-k, top-p,
-    min-p), repetition penalty, ``max_new_tokens`` / ``max_length``, EOS / pad, stop strings,
-    ``num_return_sequences``, left-padded batches and a streamer. Anything else (beam search,
-    assisted / contrastive decoding, custom logits processors or stopping criteria, scores /
-    dict outputs, ...) is transformers' own ``generate()``.
-
-    It runs bf16 models, on the weights as loaded (no quantisation). Extra keyword
-    arguments: ``fused_generate=False`` always uses transformers' path; ``fused_spec`` the
-    speculative draft length (default 6, used while it pays; 0 = off)."""
+    """``generate()`` through the fused engine when every setting is supported, else transformers'.
+    Extra kwargs: ``fused_generate``, ``fused_spec`` (draft length, 0 = off)."""
 
     def generate(self, inputs=None, generation_config=None, logits_processor=None, stopping_criteria=None,
                  prefix_allowed_tokens_fn=None, synced_gpus=None, assistant_model=None, streamer=None,
@@ -1598,23 +1472,26 @@ class _FusedGenerationMixin:
         """The fused decoder's output, or None when these settings need transformers' loop."""
         try:
             gc, model_kwargs = self._prepare_generation_config(copy.deepcopy(generation_config), **kwargs)
-        except Exception:   # another transformers version's signature: its own generate() decides
+        except Exception:
             return None
         if inputs is None:
             inputs = model_kwargs.get("input_ids")
         why = _fused_fallback_reason(self, inputs, gc, model_kwargs, logits_processor, stopping_criteria, hf)
-        self._fused_last_reason = why   # why the last generate() fell back (None: it ran fused)
+        self._fused_last_reason = why
         if why is not None:
             return None
         eos = gc.eos_token_id
+        if hasattr(eos, "tolist"):
+            eos = eos.tolist()
         if isinstance(eos, (list, tuple)):
-            eos = eos[0] if eos else None
-        pad = gc.pad_token_id if gc.pad_token_id is not None else eos
+            eos = [int(e) for e in eos] or None
+        first_eos = eos[0] if isinstance(eos, list) else eos
+        pad = gc.pad_token_id if gc.pad_token_id is not None else first_eos
         if isinstance(pad, (list, tuple)):
             pad = pad[0]
         mask = model_kwargs.get("attention_mask")
-        if mask is None:   # as transformers infers it: pad tokens (when pad != eos) are padding
-            mask = (inputs != pad).long() if pad is not None and pad != eos and bool((inputs == pad).any()) \
+        if mask is None:
+            mask = (inputs != pad).long() if pad is not None and pad not in (eos if isinstance(eos, list) else [eos]) and bool((inputs == pad).any()) \
                 else torch.ones_like(inputs)
         n_new = gc.max_new_tokens if gc.max_new_tokens is not None else (gc.max_length or 20) - inputs.shape[1]
         stop = gc.stop_strings or []
@@ -1626,23 +1503,39 @@ class _FusedGenerationMixin:
                 sample=bool(gc.do_sample), temperature=gc.temperature or 1.0, top_k=gc.top_k or 0,
                 top_p=1.0 if gc.top_p is None else gc.top_p, min_p=gc.min_p or 0.0,
                 repetition_penalty=gc.repetition_penalty or 1.0, num_samples=gc.num_return_sequences or 1,
-                eos=eos, pad=pad, stop=stop, tokenizer=model_kwargs.get("tokenizer"), streamer=streamer)
+                eos=eos, pad=pad, stop=stop, tokenizer=model_kwargs.get("tokenizer"), streamer=streamer,
+                stopping_criteria=stopping_criteria or None)
 
     def _fused_decoder(self, spec: int):
         """The decoder of these weights, kept between calls (rebuilt when a parameter changed)."""
         from .luau_infer import FusedDecoder
         try:
             key = (spec,) + tuple((p.data_ptr(), p._version) for p in self.parameters())
-        except RuntimeError:   # inference tensors have no version counter
+        except RuntimeError:
             key = (spec,) + tuple(p.data_ptr() for p in self.parameters())
         hit = self.__dict__.get("_fused")
         if hit is None or hit[0] != key:
-            self.__dict__["_fused"] = None   # free the old engine's caches first
+            self.__dict__["_fused"] = None
             with torch.inference_mode():
                 self.__dict__["_fused"] = (key, FusedDecoder(self, None, spec))
         return self.__dict__["_fused"][1]
 
 
+class LuauCoderHeadGate(nn.Module):
+    """Gated low-rank head term: input ``[h, sigmoid(u . h + b) * P h]``, ``lm_head.weight = [W, Z]``."""
+
+    def __init__(self, hidden_size: int, rank: int):
+        super().__init__()
+        self.rank = rank
+        self.proj = nn.Parameter(torch.empty(rank + 1, hidden_size))
+        self.bias = nn.Parameter(torch.zeros(1))
+
+    def forward(self, h: Tensor) -> Tensor:
+        y = F.linear(h, self.proj.to(h.dtype))
+        gate = torch.sigmoid(y[..., self.rank:] + self.bias.to(h.dtype))
+        return torch.cat([h, gate * y[..., :self.rank]], dim=-1)
+
+
 class LuauCoderForCausalLM(LuauCoderPreTrainedModel, _FusedGenerationMixin, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed.weight"} if _V5 else ["lm_head.weight"]
 
@@ -1650,7 +1543,11 @@ class LuauCoderForCausalLM(LuauCoderPreTrainedModel, _FusedGenerationMixin, Gene
         super().__init__(config)
         self.model = LuauCoderModel(config)
         self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        r = getattr(config, "head_gate_rank", 0)
+        if r and config.tie_word_embeddings:
+            raise ValueError("head_gate_rank needs an untied head")
+        self.head_gate = LuauCoderHeadGate(config.hidden_size, r) if r else None
+        self.lm_head = nn.Linear(config.hidden_size + r, config.vocab_size, bias=False)
         self.post_init()
 
     def get_input_embeddings(self) -> nn.Embedding:
@@ -1671,15 +1568,18 @@ class LuauCoderForCausalLM(LuauCoderPreTrainedModel, _FusedGenerationMixin, Gene
                 use_cache: bool | None = None, logits_to_keep: int | Tensor = 0,
                 output_hidden_states: bool | None = None, return_dict: bool | None = None,
                 **kwargs) -> CausalLMOutputWithPast:
-        """``labels`` follow the HF convention (unshifted, ``-100`` ignored);
-        ``shift_labels`` (already shifted) takes precedence. With labels and
-        grad enabled the fused LM-head + CE kernel runs and ``logits`` is None."""
+        """HF-convention ``labels``; with labels and grad the fused CE runs and ``logits`` is None."""
         num_items_in_batch = kwargs.pop("num_items_in_batch", None)
         shift_labels = kwargs.pop("shift_labels", None)
         out = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
                          past_key_values=past_key_values, inputs_embeds=inputs_embeds,
                          use_cache=use_cache, output_hidden_states=output_hidden_states, **kwargs)
         h = out.last_hidden_state
+        if self.head_gate is not None:
+            if shift_labels is None and labels is None and isinstance(logits_to_keep, int) and logits_to_keep:
+                h = h[:, -logits_to_keep:, :]
+                logits_to_keep = 0
+            h = self.head_gate(h)
         if labels is not None and shift_labels is None:
             shift_labels = F.pad(labels[:, 1:], (0, 1), value=IGNORE_INDEX)
         loss = logits = None