Retire the optional kernels, carry the tokenizer, drop dead config keys

added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<|rwkv_tokenizer_end_of_text|>": 0
+}

config.json

@@ -9,13 +9,10 @@
   },
   "bos_token_id": 0,
   "decay_low_rank_dim": 64,
-  "deep_embed_size": 768,
   "dtype": "bfloat16",
   "eos_token_id": 0,
   "gate_low_rank_dim": 128,
   "head_dim": 64,
-  "hidden_act": "sqrelu",
-  "hidden_ratio": 4.0,
   "hidden_size": 768,
   "intermediate_size": 3072,
   "max_position_embeddings": 4096,
@@ -25,7 +22,6 @@
   "num_heads": 12,
   "num_hidden_layers": 12,
   "pad_token_id": 0,
-  "sparse_channel_mix": false,
   "tie_word_embeddings": false,
   "transformers_version": "5.15.0.dev0",
   "use_cache": true,

configuration_rwkv7.py

@@ -13,107 +13,88 @@
 # limitations under the License.
 """RWKV-7 (Goose) model configuration."""
 
+from huggingface_hub.dataclasses import strict
+
 from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import auto_docstring
 
 
-class Rwkv7Config(PreTrainedConfig):
-    r"""
+@auto_docstring(
+    custom_intro="""
     Configuration for [`Rwkv7Model`], an all-recurrent (attention-free) RWKV-7 "Goose"
     model. Instantiating with the defaults yields the ~0.1B RWKV-7 configuration.
 
     Parameter names follow the upstream RWKV reference implementation
     (`BlinkDL/RWKV-LM`) rather than a renamed variant, so converting a native
     `.pth` checkpoint is close to a rename-free copy.
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 65536):
-            Vocabulary size (RWKV "world" tokenizer).
-        hidden_size (`int`, *optional*, defaults to 768):
-            Model width `C`.
-        num_hidden_layers (`int`, *optional*, defaults to 12):
-            Number of blocks.
-        head_dim (`int`, *optional*, defaults to 64):
-            Width of one WKV head. `hidden_size` must be divisible by it.
-        num_heads (`int`, *optional*, defaults to 12):
-            Number of WKV heads; must equal `hidden_size // head_dim`.
-        decay_low_rank_dim (`int`, *optional*, defaults to 64):
-            Rank of the decay (`w`) LoRA.
-        a_low_rank_dim (`int`, *optional*, defaults to 64):
-            Rank of the in-context-learning-rate (`a`) LoRA.
-        v_low_rank_dim (`int`, *optional*, defaults to 32):
-            Rank of the value-residual (`v`) LoRA. Unused on layer 0, which
-            *produces* `v_first` instead of mixing towards it.
-        gate_low_rank_dim (`int`, *optional*, defaults to 128):
-            Rank of the output-gate (`g`) LoRA.
-        intermediate_size (`int`, *optional*):
-            Channel-mix inner width (`hidden_ratio * hidden_size` by convention).
-        hidden_ratio (`float`, *optional*, defaults to 4.0):
-            Used to derive `intermediate_size` when it is not given explicitly.
-        hidden_act (`str`, *optional*, defaults to `"sqrelu"`):
-            Channel-mix activation. RWKV-7 uses squared ReLU.
-        norm_eps (`float`, *optional*, defaults to 1e-05):
-            Epsilon of every LayerNorm/GroupNorm in the model.
-        norm_bias (`bool`, *optional*, defaults to `True`):
-            Whether the norms carry a bias.
-        max_position_embeddings (`int`, *optional*, defaults to 8192):
-            Training context length. RWKV is recurrent and not bounded by it at
-            inference; it only sizes generation defaults.
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
-            Whether to tie the input embedding and the LM head.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether to return the recurrent state.
-        use_deep_embed (`bool`, *optional*, defaults to `False`):
-            Enable the RWKV-8 "DeepEmbed" hook: a per-layer, per-token vector that
-            channelwise-modulates the channel-mix. The table is deliberately NOT a
-            model weight — it is meant to live in RAM/SSD and be prefetched per
-            token, which is the whole point of the design (VRAM savings) — so it is
-            passed to the forward as `deep_embeds` instead. No RWKV-7 checkpoint
-            carries one; this is an extension point, off by default.
-        deep_embed_size (`int`, *optional*):
-            Width of one layer's DeepEmbed vector. `hidden_size` reproduces the
-            reference "1x" variant (modulating the channel-mix output);
-            `intermediate_size` reproduces "4x" (modulating its input). Defaults to
-            `hidden_size` when `use_deep_embed` is set.
-        wkv_state_dtype (`str`, *optional*, defaults to `"float32"`):
-            Precision the recurrent WKV state is carried and accumulated in,
-            independently of the activation dtype. The recurrence is unrolled over
-            the whole sequence, so a narrow state drifts; `"float32"` with fp16
-            activations is the combination the reference implementation uses.
-            `"float16"`/`"bfloat16"` trade that for a smaller state.
-        wkv_implementation (`str`, *optional*, defaults to `"eager"`):
-            Which WKV recurrence to use, by name, from
-            `models.rwkv7.modeling_rwkv7.RWKV7_WKV_FUNCTIONS`. `"eager"` is the
-            portable PyTorch path — the sequential step when decoding, the
-            chunk-parallel form otherwise, and per-segment when a packed batch is
-            passed. Register an entry in that mapping to plug in a fused or varlen
-            kernel without forking the model.
-        sparse_channel_mix (`bool`, *optional*, defaults to `False`):
-            Skip the channel-mix value-projection rows whose input channel is zero.
-            The activation is a squared ReLU, so its zeros are exact and skipping
-            them is exact too; on a 7.2B checkpoint only about a tenth of the
-            channels are nonzero -- 10.07% measured over 16 decode steps, 6.85%
-            low and 11.99% high -- and that projection is a third of the
-            model's bytes. Costs a
-            transposed copy of the value weight (about +30% weights), built lazily,
-            and only pays once launches are captured — see the model doc.
-
-            Exact but not bit-reproducible: the surviving channels are summed across
-            several partitions that combine through an atomic add, so the order the
-            partitions arrive in — and therefore the last unit in the last place of
-            each output — varies between runs of identical input. Same rounding
-            class as a split-K GEMM. Greedy decoding can turn that into a different
-            token at a near-tie, so leave this off when you need a run to reproduce
-            itself exactly.
-        bos_token_id (`int`, *optional*, defaults to 0):
-            Beginning-of-sequence id. The RWKV world tokenizer has no dedicated BOS
-            token and the reference implementation prepends nothing, so this exists to
-            satisfy `GenerationMixin` rather than to be emitted.
-        eos_token_id (`int`, *optional*, defaults to 0):
-            End-of-sequence id, id 0 in the RWKV world vocabulary.
-        pad_token_id (`int`, *optional*, defaults to 0):
-            Padding id, the same id 0. Set deliberately rather than left `None`:
-            `generate` needs one to pad a batch, and without it a batched call either
-            raised or fell back to the eos id with a warning on every step.
+    """,
+    checkpoint="Hakureirm/rwkv7-168m-pile-hf",
+)
+@strict
+class Rwkv7Config(PreTrainedConfig):
+    r"""
+    vocab_size (`int`, *optional*, defaults to 65536):
+        Vocabulary size (RWKV "world" tokenizer).
+    hidden_size (`int`, *optional*, defaults to 768):
+        Model width `C`.
+    num_hidden_layers (`int`, *optional*, defaults to 12):
+        Number of blocks.
+    head_dim (`int`, *optional*, defaults to 64):
+        Width of one WKV head. `hidden_size` must be divisible by it.
+    num_heads (`int`, *optional*, defaults to 12):
+        Number of WKV heads; must equal `hidden_size // head_dim`.
+    decay_low_rank_dim (`int`, *optional*, defaults to 64):
+        Rank of the decay (`w`) LoRA.
+    a_low_rank_dim (`int`, *optional*, defaults to 64):
+        Rank of the in-context-learning-rate (`a`) LoRA.
+    v_low_rank_dim (`int`, *optional*, defaults to 32):
+        Rank of the value-residual (`v`) LoRA. Unused on layer 0, which
+        *produces* `v_first` instead of mixing towards it.
+    gate_low_rank_dim (`int`, *optional*, defaults to 128):
+        Rank of the output-gate (`g`) LoRA.
+    intermediate_size (`int`, *optional*):
+        Channel-mix inner width. Defaults to `4 * hidden_size`.
+    norm_eps (`float`, *optional*, defaults to 1e-05):
+        Epsilon of every LayerNorm/GroupNorm in the model.
+    norm_bias (`bool`, *optional*, defaults to `True`):
+        Whether the norms carry a bias.
+    max_position_embeddings (`int`, *optional*, defaults to 8192):
+        Training context length. RWKV is recurrent and not bounded by it at
+        inference; it only sizes generation defaults.
+    tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+        Whether to tie the input embedding and the LM head.
+    use_cache (`bool`, *optional*, defaults to `True`):
+        Whether to return the recurrent state.
+    use_deep_embed (`bool`, *optional*, defaults to `False`):
+        Enable the RWKV-8 "DeepEmbed" hook: a per-layer, per-token vector that
+        channelwise-modulates the channel-mix. The table is deliberately NOT a
+        model weight. It is meant to live in RAM/SSD and be prefetched per
+        token, which is the whole point of the design (VRAM savings), so it is
+        passed to the forward as `deep_embeds` instead. No RWKV-7 checkpoint
+        carries one; this is an extension point, off by default.
+    wkv_state_dtype (`str`, *optional*, defaults to `"float32"`):
+        Precision the recurrent WKV state is carried and accumulated in,
+        independently of the activation dtype. The recurrence is unrolled over
+        the whole sequence, so a narrow state drifts; `"float32"` with fp16
+        activations is the combination the reference implementation uses.
+        `"float16"`/`"bfloat16"` trade that for a smaller state.
+    wkv_implementation (`str`, *optional*, defaults to `"eager"`):
+        Which WKV recurrence to use, by name, from
+        `models.rwkv7.modeling_rwkv7.RWKV7_WKV_FUNCTIONS`. `"eager"` is the
+        portable PyTorch path: the sequential step when decoding, the
+        chunk-parallel form otherwise, and per-segment when a packed batch is
+        passed. Register an entry in that mapping to plug in a fused or varlen
+        kernel without forking the model.
+    bos_token_id (`int`, *optional*, defaults to 0):
+        Beginning-of-sequence id. The RWKV world tokenizer has no dedicated BOS
+        token and the reference implementation prepends nothing, so this exists to
+        satisfy `GenerationMixin` rather than to be emitted.
+    eos_token_id (`int`, *optional*, defaults to 0):
+        End-of-sequence id, id 0 in the RWKV world vocabulary.
+    pad_token_id (`int`, *optional*, defaults to 0):
+        Padding id, the same id 0. Set deliberately rather than left `None`:
+        `generate` needs one to pad a batch, and without it a batched call either
+        raised or fell back to the eos id with a warning on every step.
 
     ```python
     >>> from transformers import Rwkv7Config, Rwkv7Model
@@ -126,81 +107,46 @@ class Rwkv7Config(PreTrainedConfig):
     model_type = "rwkv7"
     keys_to_ignore_at_inference = ["past_key_values"]
 
-    def __init__(
-        self,
-        vocab_size=65536,
-        hidden_size=768,
-        num_hidden_layers=12,
-        head_dim=64,
-        num_heads=12,
-        decay_low_rank_dim=64,
-        a_low_rank_dim=64,
-        v_low_rank_dim=32,
-        gate_low_rank_dim=128,
-        intermediate_size=None,
-        hidden_ratio=4.0,
-        hidden_act="sqrelu",
-        norm_eps=1e-5,
-        norm_bias=True,
-        max_position_embeddings=8192,
-        tie_word_embeddings=False,
-        use_cache=True,
-        use_deep_embed=False,
-        deep_embed_size=None,
-        wkv_state_dtype="float32",
-        wkv_implementation="eager",
-        sparse_channel_mix=False,
-        bos_token_id=0,
-        eos_token_id=0,
-        pad_token_id=0,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.head_dim = head_dim
-        self.num_heads = num_heads
-        self.decay_low_rank_dim = decay_low_rank_dim
-        self.a_low_rank_dim = a_low_rank_dim
-        self.v_low_rank_dim = v_low_rank_dim
-        self.gate_low_rank_dim = gate_low_rank_dim
-        self.hidden_ratio = hidden_ratio
-        # `None` rather than a number, so `hidden_ratio` is what actually decides. It
-        # used to default to 3072, which is exactly `768 * 4.0` -- correct for the
-        # default `hidden_size` and silently wrong for every other one. A config built
-        # as `Rwkv7Config(hidden_size=4096, num_heads=64)` came back with a channel-mix
-        # four times narrower than the architecture it names, and `hidden_ratio` was
-        # dead code that nothing could reach without passing `intermediate_size=None`
-        # explicitly. Serialisation is unaffected: the resolved value is written to
-        # `config.json` either way.
-        self.intermediate_size = (
-            intermediate_size if intermediate_size is not None else int(hidden_size * hidden_ratio)
-        )
-        self.hidden_act = hidden_act
-        self.norm_eps = norm_eps
-        self.norm_bias = norm_bias
-        self.max_position_embeddings = max_position_embeddings
-        self.use_cache = use_cache
-        self.use_deep_embed = use_deep_embed
-        self.deep_embed_size = deep_embed_size if deep_embed_size is not None else hidden_size
-        if wkv_state_dtype not in ("float32", "float16", "bfloat16"):
-            raise ValueError(f"wkv_state_dtype must be float32/float16/bfloat16, got {wkv_state_dtype}")
-        self.wkv_state_dtype = wkv_state_dtype
-        self.wkv_implementation = wkv_implementation
-        self.sparse_channel_mix = sparse_channel_mix
-
-        if hidden_size % head_dim != 0:
-            raise ValueError(f"hidden_size {hidden_size} must be divisible by head_dim {head_dim}")
-        if num_heads != hidden_size // head_dim:
-            raise ValueError(f"num_heads must be hidden_size // head_dim = {hidden_size // head_dim}, got {num_heads}")
+    vocab_size: int = 65536
+    hidden_size: int = 768
+    num_hidden_layers: int = 12
+    head_dim: int = 64
+    num_heads: int = 12
+    decay_low_rank_dim: int = 64
+    a_low_rank_dim: int = 64
+    v_low_rank_dim: int = 32
+    gate_low_rank_dim: int = 128
+    # `None` rather than a number: a literal default is correct for the default
+    # `hidden_size` and silently wrong for every other one, so a config built as
+    # `Rwkv7Config(hidden_size=4096, num_heads=64)` would come back with a channel-mix
+    # four times narrower than the architecture it names. `__post_init__` resolves it,
+    # and the resolved value is written to `config.json` either way.
+    intermediate_size: int | None = None
+    norm_eps: float = 1e-5
+    norm_bias: bool = True
+    max_position_embeddings: int = 8192
+    tie_word_embeddings: bool = False
+    use_cache: bool = True
+    use_deep_embed: bool = False
+    wkv_state_dtype: str = "float32"
+    wkv_implementation: str = "eager"
+    bos_token_id: int | None = 0
+    eos_token_id: int | None = 0
+    pad_token_id: int | None = 0
 
-        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,
-        )
+    def __post_init__(self, **kwargs):
+        if self.intermediate_size is None:
+            self.intermediate_size = 4 * self.hidden_size
+        if self.wkv_state_dtype not in ("float32", "float16", "bfloat16"):
+            raise ValueError(f"wkv_state_dtype must be float32/float16/bfloat16, got {self.wkv_state_dtype}")
+        if self.hidden_size % self.head_dim != 0:
+            raise ValueError(f"hidden_size {self.hidden_size} must be divisible by head_dim {self.head_dim}")
+        if self.num_heads != self.hidden_size // self.head_dim:
+            raise ValueError(
+                f"num_heads must be hidden_size // head_dim = {self.hidden_size // self.head_dim}, "
+                f"got {self.num_heads}"
+            )
+        super().__post_init__(**kwargs)
 
 
 __all__ = ["Rwkv7Config"]

fused_wkv.py

@@ -1,138 +0,0 @@
-# Copyright 2026 The RWKV team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Single-token WKV in one Triton kernel: the state is read once and written once.
-
-The portable recurrence in `modeling_rwkv7.rwkv7_recurrent` touches the state four
-times for one decoded token -- once to form `sa = (-kk) @ S`, once to update `S`,
-once to read it back for `r @ S`, once to store it. That is the natural way to write
-it and it is free at batch 1, where the state is a megabyte a layer and invisible
-next to the weights being streamed.
-
-It stops being free with batch. At batch 256 and 7.2B's shape the state is
-`[256, 64, 64, 64]` -- 268 MB a layer, 8.6 GB across 32 layers -- and profiling put
-77% of a decode step inside the time-mix while its own projections were 5.5% of it.
-The cost was the state, not the weights, and four passes over it rather than two.
-
-This does the whole update with the `[head_dim, head_dim]` tile resident. Measured on
-one RTX 5090 at 7.2B: 5.7x on the recurrence at batch 256, and the whole model's
-batched decode goes from 55% of the albatross reference to 85%.
-
-Called as a plain function rather than wrapped in a `custom_op`. Dynamo traces a
-`@triton.jit` launch natively, and the wrapper is opaque to it -- with the state
-declared as mutated, inductor stops issuing CUDA graphs for the region. That is
-invisible at batch 256, where the kernel wins anyway, and costs everything at batch
-1: 138.9 tok/s wrapped becomes 57.9, which is eager speed. `sparse_channel_mix` is a
-`custom_op` because its three kernels and its compaction buffers are genuinely
-opaque; one traced kernel is better off without the wrapper.
-"""
-
-import torch
-from transformers.utils.import_utils import is_triton_available
-
-# Guarded rather than bare, although the module itself is only ever imported
-# under the same `is_triton_available()` condition: import scanners that parse
-# this file in isolation (transformers' remote-code `check_imports` is one)
-# exempt imports inside an availability-guarded block, and demand triton on
-# every machine otherwise.
-if is_triton_available():
-    import triton
-    import triton.language as tl
-
-
-@triton.jit
-def _wkv_one_kernel(
-    state_ptr,
-    r_ptr,
-    w_ptr,
-    k_ptr,
-    v_ptr,
-    kk_ptr,
-    a_ptr,
-    out_ptr,
-    stride_sb,
-    stride_sh,
-    stride_si,
-    stride_sj,
-    stride_vb,
-    stride_vh,
-    N: tl.constexpr,
-):
-    """One program per (batch, head).
-
-    `N` is `head_dim` and is a compile-time constant, so the `[N, N]` tile lives in
-    registers for the whole body: 64x64 fp32 is 16 KB, which is what makes reading
-    the state once sufficient.
-    """
-    batch = tl.program_id(0)
-    head = tl.program_id(1)
-
-    i = tl.arange(0, N)
-    j = tl.arange(0, N)
-    vec = batch * stride_vb + head * stride_vh
-    r = tl.load(r_ptr + vec + i).to(tl.float32)
-    w = tl.load(w_ptr + vec + i).to(tl.float32)
-    k = tl.load(k_ptr + vec + i).to(tl.float32)
-    v = tl.load(v_ptr + vec + j).to(tl.float32)
-    kk = tl.load(kk_ptr + vec + i).to(tl.float32)
-    a = tl.load(a_ptr + vec + i).to(tl.float32)
-
-    # Every axis of the state is strided, the last one included. Hardcoding unit
-    # stride there reads the tile transposed when the caller hands over anything
-    # that is not contiguous, and the result is wrong by 80-98% with nothing
-    # raised: the shape is right, so neither the kernel nor Triton notices.
-    off = batch * stride_sb + head * stride_sh + i[:, None] * stride_si + j[None, :] * stride_sj
-    state = tl.load(state_ptr + off).to(tl.float32)
-
-    # Same three lines as the reference, in the same order and in fp32: `sa` uses the
-    # PRE-update state and the output uses the POST-update one.
-    sa = tl.sum((-kk)[:, None] * state, axis=0)
-    state = tl.exp(w)[:, None] * state + (kk * a)[:, None] * sa[None, :] + k[:, None] * v[None, :]
-    out = tl.sum(r[:, None] * state, axis=0)
-
-    tl.store(state_ptr + off, state.to(state_ptr.dtype.element_ty))
-    tl.store(out_ptr + vec + j, out.to(out_ptr.dtype.element_ty))
-
-
-def fused_wkv_one(
-    r: torch.Tensor,
-    w_log: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    kk: torch.Tensor,
-    a: torch.Tensor,
-    state: torch.Tensor,
-) -> torch.Tensor:
-    batch, _, heads, head_dim = r.shape
-    # The grid comes from `r`, so a state that disagrees with it is read at offsets
-    # belonging to another row and returns a plausible answer for the wrong sequence.
-    # Checked rather than assumed: the kernel cannot tell, and neither can the caller.
-    if tuple(state.shape) != (batch, heads, head_dim, head_dim):
-        raise ValueError(
-            f"state has shape {tuple(state.shape)}, but the vectors imply {(batch, heads, head_dim, head_dim)}"
-        )
-    out = torch.empty(batch, heads, head_dim, device=r.device, dtype=r.dtype)
-    flat = [t.reshape(batch, heads, head_dim).contiguous() for t in (r, w_log, k, v, kk, a)]
-    _wkv_one_kernel[(batch, heads)](
-        state,
-        *flat,
-        out,
-        state.stride(0),
-        state.stride(1),
-        state.stride(2),
-        state.stride(3),
-        out.stride(0),
-        out.stride(1),
-        N=head_dim,
-    )
-    return out.view(batch, 1, heads, head_dim)

hf_rwkv_tokenizer.py

@@ -0,0 +1,279 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for RWKV."""
+
+import os
+import re
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+from transformers.utils import logging
+
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.get_logger(__name__)
+
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "rwkv_vocab_v20230424.txt",
+}
+
+class TRIE:
+    __slots__ = tuple("ch,to,values,front".split(","))
+    to: list
+    values: set
+
+    def __init__(self, front=None, ch=None):
+        self.ch = ch
+        self.to = [None for ch in range(256)]
+        self.values = set()
+        self.front = front
+
+    def __repr__(self):
+        fr = self
+        ret = []
+        while fr != None:
+            if fr.ch != None:
+                ret.append(fr.ch)
+            fr = fr.front
+        return "<TRIE %s %s>" % (ret[::-1], self.values)
+
+    def add(self, key: bytes, idx: int = 0, val=None):
+        if idx == len(key):
+            if val is None:
+                val = key
+            self.values.add(val)
+            return self
+        ch = key[idx]
+        if self.to[ch] is None:
+            self.to[ch] = TRIE(front=self, ch=ch)
+        return self.to[ch].add(key, idx=idx + 1, val=val)
+
+    def find_longest(self, key: bytes, idx: int = 0):
+        u: TRIE = self
+        ch: int = key[idx]
+
+        while u.to[ch] is not None:
+            u = u.to[ch]
+            idx += 1
+            if u.values:
+                ret = idx, u, u.values
+            if idx == len(key):
+                break
+            ch = key[idx]
+        return ret
+
+
+class RWKV_TOKENIZER:
+    def __init__(self, file_name):
+        self.idx2token = {}
+        sorted = []  # must be already sorted
+        with open(file_name, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        for l in lines:
+            idx = int(l[: l.index(" ")])
+            x = eval(l[l.index(" ") : l.rindex(" ")])
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+
+            assert len(x) == int(l[l.rindex(" ") :])
+            sorted += [x]
+            self.idx2token[idx] = x
+
+        self.token2idx = {}
+        for k, v in self.idx2token.items():
+            self.token2idx[v] = int(k)
+
+        self.root = TRIE()
+        for t, i in self.token2idx.items():
+            _ = self.root.add(t, val=(t, i))
+
+    def encodeBytes(self, src: bytes):
+        idx: int = 0
+        tokens = []
+        while idx < len(src):
+            _idx: int = idx
+            idx, _, values = self.root.find_longest(src, idx)
+            assert idx != _idx
+            _, token = next(iter(values))
+            tokens.append(token)
+        return tokens
+
+    def decodeBytes(self, tokens):
+        return b"".join(map(lambda i: self.idx2token[i], tokens))
+
+    def encode(self, src):
+        if isinstance(src, str):
+            return [self.encodeBytes(src.encode("utf-8"))]
+        elif isinstance(src, list):
+            return [self.encodeBytes(s.encode("utf-8")) for s in src]
+
+    def decode(self, tokens):
+        return [self.decodeBytes(batch).decode("utf-8") for batch in tokens]
+        # try:
+        #     return self.decodeBytes(tokens).decode('utf-8')
+        # except:
+        #     return '\ufffd' # bad utf-8
+
+    def printTokens(self, tokens):
+        for i in tokens:
+            s = self.idx2token[i]
+            try:
+                s = s.decode("utf-8")
+            except:
+                pass
+            print(f"{repr(s)}{i}", end=" ")
+        print()
+
+
+class RwkvTokenizer(PreTrainedTokenizer):
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self, vocab_file, bos_token="<|rwkv_tokenizer_end_of_text|>", eos_token="<|rwkv_tokenizer_end_of_text|>", unk_token="<|rwkv_tokenizer_end_of_text|>", **kwargs
+    ):
+        if not os.path.isfile(vocab_file):
+            raise ValueError(
+                f"Can't find a vocabulary file at path '{vocab_file}'."
+            )
+
+        with open(vocab_file, "r", encoding="utf-8") as reader:
+            tokens = reader.readlines()
+
+        if "add_bos_token" in kwargs:
+            self.add_bos_token = kwargs["add_bos_token"]
+        else:
+            self.add_bos_token = False
+        self.trie_tokenizer = RWKV_TOKENIZER(vocab_file)
+        vocab = self.trie_tokenizer.token2idx
+        self.encoder = vocab
+        self.decoder = {v: k for k, v in vocab.items()}
+        self._added_tokens_decoder = {0: AddedToken(str(bos_token))}
+        super().__init__(
+            bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, **kwargs
+        )
+
+    @property
+    def vocab_size(self):
+        return len(self.encoder)
+
+    def get_vocab(self):
+        vocab = self.encoder
+        vocab.update(self.added_tokens_encoder)
+        vocab = dict(sorted(vocab.items(), key=lambda item: item[1]))
+        return vocab
+
+    def _tokenize(self, text, split_special_tokens=False):
+        # return self.wordpiece_tokenizer.tokenize(text.encode("utf-8"))
+        return self.trie_tokenizer.encode(text)[0]
+
+    def _convert_token_to_id(self, token):
+        return token
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (byte) using the vocab."""
+        token = self.decoder.get(index, self.unk_token)
+        if isinstance(token, (bytes)):
+            token = token.decode("utf-8", errors="replace")
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (bytes) in a single string. Additional tokens are encoded to bytes"""
+        out_string = b"".join(
+            [k.encode(errors="replace") if isinstance(k, str) else k for k in tokens]
+        ).decode("utf-8")
+        return out_string
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        index = 0
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory,
+                (filename_prefix + "-" if filename_prefix else "") + "vocab.txt",
+            )
+        else:
+            vocab_file = (
+                filename_prefix + "-" if filename_prefix else ""
+            ) + save_directory
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(
+                self.encoder.items(), key=lambda kv: kv[1]
+            ):
+                if index != token_index:
+                    logger.warning(
+                        f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
+                        " Please check that the vocabulary is not corrupted!"
+                    )
+                    index = token_index
+                writer.write(str(token) + "\n")
+                index += 1
+        return (vocab_file,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        if self.add_bos_token:
+            bos_token_ids = [self.bos_token_id]
+        else:
+            bos_token_ids = []
+
+        output = bos_token_ids + token_ids_0
+
+        if token_ids_1 is None:
+            return output
+
+        return output + bos_token_ids + token_ids_1
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        """
+        Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True,
+            )
+
+        if not self.add_bos_token:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=False,
+            )
+
+        if token_ids_1 is None:
+            return [1] + ([0] * len(token_ids_0))
+        return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))

modeling_rwkv7.py

@@ -30,22 +30,9 @@ from transformers.initialization import zeros_
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import ModelOutput, auto_docstring, can_return_tuple, logging
-from transformers.utils.import_utils import is_triton_available
 from .configuration_rwkv7 import Rwkv7Config
 
 
-# Imported at module scope, NOT lazily from the forward. A function containing an
-# `import` is on dynamo's skip list whether or not the import actually runs, so
-# calling one per layer breaks the graph at every layer -- which cost 138.9 tok/s at
-# batch 1 down to 47.4, eager speed, while the kernel itself was fine. `sparse_channel_mix`
-# avoids this by importing from `allocate_state`, outside anything traced; there is no
-# equivalent hook here, so it happens once at import.
-if is_triton_available():
-    from .fused_wkv import fused_wkv_one
-else:  # pragma: no cover - exercised only where triton is absent
-    fused_wkv_one = None
-
-
 logger = logging.get_logger(__name__)
 
 
@@ -72,7 +59,7 @@ class Rwkv7TokenShift(nn.Module):
         if cu_seq_lens is not None:
             # In a packed row the token before a segment's first one belongs to the
             # PREVIOUS sequence. Resetting the recurrent state per segment does not
-            # cover this — the shift reaches back through it — so the first token of
+            # cover this (the shift reaches back through it), so the first token of
             # each segment gets the zero shift a sequence start is supposed to see.
             positions = torch.arange(x.shape[1], device=x.device)
             starts = (positions[:, None] == cu_seq_lens[None, :-1]).any(dim=1)
@@ -97,7 +84,7 @@ class Rwkv7TokenShift(nn.Module):
 
 # e^-0.5. The decay LoRA emits w_log = -INV_SQRT_E * sigmoid(...), so w_log lies in
 # (-e^-0.5, 0) and the per-step decay exp(w_log) lies in (exp(-e^-0.5), 1), i.e.
-# (0.5452, 1) — see RWKV-7 reference. Note the floor is exp(-e^-0.5), not e^-0.5: this
+# (0.5452, 1). See RWKV-7 reference. Note the floor is exp(-e^-0.5), not e^-0.5: this
 # comment said the latter, and `rwkv7_chunked` built its chunk_size bound on it.
 _INV_SQRT_E = 0.6065306597126334
 
@@ -126,7 +113,7 @@ def rwkv7_recurrent(
         out   = r @ S                     # uses the POST-update state
 
     The state axes are (key, value); `S[i, j]` accumulates key channel `i` against
-    value channel `j`. Accumulation is fp32 regardless of the activation dtype —
+    value channel `j`. Accumulation is fp32 regardless of the activation dtype:
     the recurrence is unrolled over the whole sequence, so a lower-precision state
     drifts. This is the portable path; a fused kernel may replace it as long as it
     reproduces these values.
@@ -180,14 +167,14 @@ def rwkv7_chunked(
 
     `chunk_size` is bounded by that division by `c`, but by OVERFLOW rather than by
     precision, and the difference is worth a factor of four. The per-step decay is at
-    least `exp(-e^-0.5)` = 0.5452, so at the worst case — every channel pinned at that
-    floor — `1/c` grows like `e^(e^-0.5 * chunk_size)`. That reaches 7.2e16 at 64 and
+    least `exp(-e^-0.5)` = 0.5452, so at the worst case, every channel pinned at that
+    floor, `1/c` grows like `e^(e^-0.5 * chunk_size)`. That reaches 7.2e16 at 64 and
     fp32 tops out near 3.4e38, so there are twenty-odd decades of headroom left at the
     default, and the ceiling is `ln(finfo.max) / e^-0.5` = 146 in fp32.
 
     That derivation is worth stating carefully because the first version of it was
     wrong in a way the numbers hid. It said the decay floor was `e^-0.5` rather than
-    `exp(-e^-0.5)`, hence growth like `e^(0.5 * chunk_size)` — which puts the ceiling at
+    `exp(-e^-0.5)`, hence growth like `e^(0.5 * chunk_size)`, which puts the ceiling at
     177, and measured at the decay floor this function returns all-NaN from 147 up. The
     quoted 7.2e16 came from the correct law all along (`e^(0.5*64)` is 7.9e13), so the
     arithmetic had been done right and written up wrong, and following the prose rather
@@ -292,7 +279,7 @@ def rwkv7_eager(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Portable WKV: the sequential step for a single token, chunk-parallel otherwise.
 
-    With `cu_seq_lens` the row is a *packed* batch — several independent sequences
+    With `cu_seq_lens` the row is a *packed* batch: several independent sequences
     laid end to end, the layout a varlen kernel consumes. Each one has to start
     from a fresh state, so they are run in turn and their outputs concatenated;
     the state returned is the last segment's, which is the one a continuation
@@ -326,55 +313,6 @@ def rwkv7_eager(
 RWKV7_WKV_FUNCTIONS = {"eager": rwkv7_eager}
 
 
-def rwkv7_fused(
-    r: torch.Tensor,
-    w_log: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    kk: torch.Tensor,
-    a: torch.Tensor,
-    state: torch.Tensor,
-    cu_seq_lens: torch.Tensor | None = None,
-    **kwargs,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """`"eager"`, with the single-token step replaced by a fused Triton kernel.
-
-    Only the decode step changes; prefill and packed batches fall through to the
-    portable path, which is also where this lands when Triton is missing or the tensors
-    are on CPU. A narrower state is NOT a fallback case -- the kernel stores back
-    through the state's own dtype, which is the point of letting `wkv_state_dtype`
-    select it, and this docstring claimed the opposite for several commits after that
-    fallback was deliberately removed. The state is handed back rather than mutated in
-    place; the portable path does not mutate either, and saying it did was the same
-    stale prose.
-
-    Worth it only with batch: the kernel's whole point is that the state is read once
-    instead of three times, and at batch 1 the state is not what the step is spending
-    its time on. Measured at 7.2B on one RTX 5090, against albatross faster3a: batched
-    decode 55% -> 85% at 256x1, 91% -> 100% at 32x1, and 1x1 unchanged.
-    """
-    single_token = r.shape[1] == 1 and cu_seq_lens is None
-    # The kernel loads to fp32 and stores back in whatever the state's dtype is, so a
-    # narrower state is not a different code path -- and it is the only lever left at
-    # batch. The kernel already runs at ~86% of this card's bandwidth moving a 268 MB
-    # fp32 state per layer, so halving the state is worth more than any tuning of it.
-    # `head_dim` a power of two is part of the kernel's applicability domain, not a
-    # detail: `tl.arange` requires it, and without this term a width like 192 reached
-    # Triton and came back as a CompilationError raised from inside the kernel rather
-    # than as a fallback. The portable path has no such constraint, so falling through
-    # is both correct and what the docstring already promises for the other cases.
-    width = r.shape[-1]
-    if not (single_token and r.is_cuda and fused_wkv_one is not None and width & (width - 1) == 0):
-        return rwkv7_eager(r, w_log, k, v, kk, a, state, cu_seq_lens=cu_seq_lens, **kwargs)
-    return fused_wkv_one(r, w_log, k, v, kk, a, state), state
-
-
-# Added after the definition rather than in the literal above: the modular converter
-# is free to move the dict, and it emits it next to `rwkv7_eager` -- ahead of this
-# function, which made the generated file raise NameError on import.
-RWKV7_WKV_FUNCTIONS["fused"] = rwkv7_fused
-
-
 class Rwkv7Attention(nn.Module):
     """RWKV-7 time-mixing block (the recurrent replacement for self-attention)."""
 
@@ -467,9 +405,9 @@ class Rwkv7Attention(nn.Module):
             #
             # The decay is held at w = exp(0) = 1, so the transition is the identity.
             #
-            # `k` and `v` are zeroed EXPLICITLY. An earlier version relied on the
-            # blanked hidden state making them zero, since the projections are
-            # bias-free -- but that is only true for a pad with nothing before it.
+            # `k` and `v` are zeroed EXPLICITLY rather than relying on the blanked
+            # hidden state to make them zero: the projections are bias-free, so that
+            # only holds for a pad with nothing before it.
             # A pad that FOLLOWS a real token still receives that token's hidden
             # state through the token shift, so `delta = shifted - 0` is non-zero,
             # and so are `k` and `v`. The update term `k v^T` then entered the state
@@ -503,7 +441,7 @@ class Rwkv7Attention(nn.Module):
         def _heads(t):
             return t.view(batch, seq_len, H, N)
 
-        # Named rather than indexed: a key that is not registered used to surface as a
+        # Named rather than indexed: an unregistered key would otherwise surface as a
         # bare `KeyError: 'chunked'` from inside the forward, several frames from the
         # config field that caused it, and a caller looping over shapes would record it
         # as "this shape did not run" instead of "this model was never built". The
@@ -562,51 +500,6 @@ class Rwkv7Attention(nn.Module):
         return w_log, a, g, v_gate
 
 
-_SPARSE_OP_READY = False
-
-
-def _ensure_sparse_op() -> bool:
-    """Import the Triton kernels once, outside anything that gets traced.
-
-    A lazy `import` inside the forward is itself on dynamo's skip list, so doing
-    this on the hot path breaks the graph at every layer even when the kernel is a
-    properly registered custom op.
-    """
-    global _SPARSE_OP_READY
-    if not _SPARSE_OP_READY and is_triton_available():
-        from . import sparse_channel_mix  # noqa: F401  (registers the op)
-
-        _SPARSE_OP_READY = True
-    return _SPARSE_OP_READY
-
-
-def sparse_channel_mix_value(
-    activation: torch.Tensor,
-    weight_t: torch.Tensor,
-    accumulator: torch.Tensor,
-    index: torch.Tensor,
-    value: torch.Tensor,
-    counter: torch.Tensor,
-) -> torch.Tensor:
-    """`out = activation @ weight_t`, reading only the rows the input selects.
-
-    `activation` is `relu(key(x))**2`, so its zeros are EXACT and a zero channel
-    contributes exactly nothing — skipping its weight row is not an approximation.
-    On a 7.2B checkpoint ~90% of channels are zero at any step (10.07% density,
-    measured over 16 decode steps) while this projection is a third of the model's
-    bytes, which is why it is worth a kernel at all. `index`/`value`/`counter` are
-    the scratch the kernel compacts the surviving channels into.
-
-    Only pays when kernel launches are captured (`torch.compile` with CUDA graphs).
-    Eagerly it LOSES to one dense cuBLAS call, because the step is then bound by
-    launch overhead rather than by the weight stream — measured both ways.
-    Falls back to a dense matmul when Triton is unavailable.
-    """
-    if _SPARSE_OP_READY:
-        return torch.ops.rwkv7.sparse_channel_mix_value(activation, weight_t, accumulator, index, value, counter)
-    return torch.nn.functional.linear(activation, weight_t.t())
-
-
 class Rwkv7FeedForward(nn.Module):
     """RWKV-7 channel-mixing block: squared-ReLU over a single token shift."""
 
@@ -619,12 +512,6 @@ class Rwkv7FeedForward(nn.Module):
         self.x_k = nn.Parameter(torch.zeros(1, 1, C))
         self.key = nn.Linear(C, config.intermediate_size, bias=False)
         self.value = nn.Linear(config.intermediate_size, C, bias=False)
-        self._value_t = None  # [inter, hidden] copy for the sparse path, built lazily
-        self._value_fingerprint = None
-        self._accumulator = None
-        self._compact_index = None
-        self._compact_value = None
-        self._compact_counter = None
 
     def forward(
         self,
@@ -643,9 +530,8 @@ class Rwkv7FeedForward(nn.Module):
         # because the design keeps the table in RAM/SSD and prefetches per token.
         # Its width says which side it attaches to -- `intermediate_size` scales the
         # projection's INPUT (the reference "4x" variant), `hidden_size` its OUTPUT
-        # ("1x"). Resolved once here so both the dense and sparse projections below
-        # see the same decision; splitting it per branch is how the sparse path
-        # silently dropped the 4x variant once already.
+        # ("1x"). Resolved once here rather than per branch, so a later branch cannot
+        # drop the 4x variant silently.
         scale_output = None
         if deep_embed is not None:
             if deep_embed.shape[-1] == inner.shape[-1]:
@@ -659,93 +545,8 @@ class Rwkv7FeedForward(nn.Module):
         return out, new_shift_state
 
     def _project(self, inner: torch.Tensor) -> torch.Tensor:
-        """`value(inner)`, sparsely when that is both enabled and worthwhile.
-
-        The sparse kernel decodes one token at a time on CUDA; anything else (a
-        batch, a prefill, CPU) takes the dense projection, which is also what makes
-        the two paths comparable in tests.
-        """
-        single_token = inner.shape[0] * inner.shape[1] == 1
-        if self.config.sparse_channel_mix and single_token and inner.is_cuda:
-            return self._sparse_value(inner.reshape(-1)).view(1, 1, -1)
         return self.value(inner)
 
-    def _sparse_value(self, activation: torch.Tensor) -> torch.Tensor:
-        """Value projection over only the nonzero channels of `relu(key(x))**2`.
-
-        The transposed weight is built on first use rather than stored: the sparse
-        read needs one contiguous row per input channel, which the `nn.Linear`
-        layout does not give, and materialising it eagerly would cost every user
-        the memory whether or not they enable this.
-        """
-        if not torch.compiler.is_compiling():
-            self.build_sparse_cache()
-        return sparse_channel_mix_value(
-            activation,
-            self._value_t,
-            self._accumulator,
-            self._compact_index,
-            self._compact_value,
-            self._compact_counter,
-        )
-
-    def build_sparse_cache(self) -> None:
-        """Materialise the transposed weight and the accumulator the sparse path needs.
-
-        Called lazily from the projection, and eagerly by
-        [`Rwkv7Model.allocate_state`] — which is what makes it safe to compile. Built
-        for the first time *inside* a compiled region these buffers cannot have their
-        addresses pinned, inductor declines CUDA graphs for a region that mutates its
-        inputs, and the decode runs 2.7x slower while reporting nothing but a single
-        line of warning.
-        """
-        weight = self.value.weight
-        _ensure_sparse_op()
-        # Tie the cache to the weight's identity AND its in-place version, so a
-        # `mul_`, a fresh load, or a replaced parameter all invalidate it. A cache
-        # that silently survives a weight change is worse than no cache.
-        #
-        # One mutation this does NOT see, stated because it is the one people write:
-        # `weight.data.copy_(...)`. Going through `.data` is what detaches from the
-        # autograd version counter, so the version is unchanged, the storage is
-        # unchanged, and the fingerprint matches a weight that no longer exists.
-        # Measured on torch 2.13: `mul_` steps the version, `.data.copy_` does not,
-        # `load_state_dict` does (it copies through the parameter). Quantisation and
-        # adapter-merging code is where `.data.copy_` shows up. Call
-        # [`invalidate_sparse_cache`] after doing that -- there is no cheap way to
-        # notice from here, and the expensive way is a device sync per decoded token.
-        fingerprint = (weight.data_ptr(), weight._version, weight.dtype)
-        if self._value_t is None or self._value_fingerprint != fingerprint:
-            inter, hidden = weight.shape[1], weight.shape[0]
-            self._value_t = weight.t().contiguous()
-            self._accumulator = torch.zeros(hidden, device=weight.device, dtype=torch.float32)
-            # The projection compacts the surviving channels into these before
-            # walking them. Allocated here with everything else rather than on
-            # first use: a buffer that first appears inside a compiled region
-            # cannot be pinned, and CUDA graphs are then declined for the whole
-            # region -- which costs far more than the projection saves.
-            self._compact_index = torch.zeros(inter, device=weight.device, dtype=torch.int32)
-            self._compact_value = torch.zeros(inter, device=weight.device, dtype=torch.float32)
-            self._compact_counter = torch.zeros(1, device=weight.device, dtype=torch.int32)
-            self._value_fingerprint = fingerprint
-
-    def invalidate_sparse_cache(self) -> None:
-        """Drop the transposed weight, so the next call rebuilds it.
-
-        The escape hatch for a weight change the fingerprint cannot see -- see
-        [`build_sparse_cache`] for which those are. Rebuilding allocates, so do it
-        before compiling, not between decode steps.
-        """
-        self._value_fingerprint = None
-
-    def _load_from_state_dict(self, *args, **kwargs):
-        # Loading replaces the projection under a cache that may already be warm.
-        # The version counter happens to catch this on current torch, but that is a
-        # property of how `load_state_dict` copies rather than a promise, and the
-        # cost of being wrong is a silently stale projection.
-        super()._load_from_state_dict(*args, **kwargs)
-        self.invalidate_sparse_cache()
-
 
 class Rwkv7CacheLayer(LinearAttentionLayer):
     """One block's slice of the recurrent state: the WKV matrix and two token shifts.
@@ -1072,24 +873,18 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
     def allocate_state(self, batch: int, device=None, dtype=None) -> Rwkv7Cache:
         """A zeroed cache, plus everything else the decode path would build lazily.
 
-        Call this before compiling, and pass the result in as `state=`. Two kinds of
-        buffer have to exist by then, and for the same reason: `mark_static_address`
-        cannot run during tracing, so anything first allocated *inside* the compiled
-        region stays unpinned, and inductor declines CUDA graphs for a region that
-        mutates its inputs. Starting from `state=None` loses them to the state
-        buffers; leaving the sparse path cold loses them to one unpinned transposed
-        weight. Either way the decode runs several times slower and says so only in
-        a line of warning, so both are handled here rather than left to the caller
-        to remember.
+        Call this before compiling, and pass the result in as `state=`.
+        `mark_static_address` cannot run during tracing, so anything first allocated
+        *inside* the compiled region stays unpinned, and inductor declines CUDA graphs
+        for a region that mutates its inputs. Starting from `state=None` loses the
+        state buffers that way, and the decode then runs several times slower while
+        saying so only in a line of warning.
         """
         state = self._empty_state(
             batch,
             device if device is not None else self.emb.weight.device,
             dtype if dtype is not None else self.emb.weight.dtype,
         )
-        if self.config.sparse_channel_mix:
-            for block in self.blocks:
-                block.ffn.build_sparse_cache()
         return state
 
     def _empty_state(self, batch: int, device, dtype) -> Rwkv7Cache:
@@ -1125,11 +920,11 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
             *history* when `cu_seq_lens` is given -- see there.
         deep_embeds (`torch.FloatTensor`, *optional*):
             RWKV-8 DeepEmbed vectors for this batch, shaped
-            `[num_layers, batch, seq_len, deep_embed_size]` (or broadcastable).
+            `[num_layers, batch, seq_len, hidden_size or intermediate_size]` (or broadcastable).
             Only meaningful when `config.use_deep_embed` is set; the table itself is
             external to the checkpoint by design.
         cu_seq_lens (`torch.LongTensor`, *optional*):
-            Cumulative sequence lengths for a *packed* batch — several sequences
+            Cumulative sequence lengths for a *packed* batch: several sequences
             concatenated into one row instead of padded to a rectangle, starting at
             0 and ending at `seq_len`, and non-decreasing. Each segment then decodes
             from a fresh recurrent state, as if it had been run on its own. This is
@@ -1178,9 +973,9 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
             # can hand over the whole conversation's mask. A prefix chunk must slice
             # its own mask to match, or it silently reads the wrong positions.
             #
-            # Applied at seq_len == 1 as well. It used to be skipped there, on the
-            # theory that a single decoded token is never padding -- true of
-            # `generate`, but an assumption about the caller rather than a property
+            # Applied at seq_len == 1 as well. Skipping it there assumes a single
+            # decoded token is never padding -- true of `generate`, but an
+            # assumption about the caller rather than a property
             # of the model, and a fully-masked 1-token row was moving the state.
             keep = attention_mask[:, -inputs_embeds.shape[1] :, None].to(inputs_embeds.dtype)
 
@@ -1217,8 +1012,8 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
             # Gated on the config flag as well as on the argument: passing
-            # `deep_embeds` to a model configured without them used to modulate the
-            # channel-mix anyway, which is a silently different model.
+            # `deep_embeds` to a model configured without them would otherwise modulate
+            # the channel-mix anyway, which is a silently different model.
             layer_deep_embed = (
                 deep_embeds[block.layer_id] if deep_embeds is not None and self.config.use_deep_embed else None
             )
@@ -1258,9 +1053,9 @@ class Rwkv7ForCausalLM(Rwkv7PreTrainedModel, GenerationMixin):
         )
         model_inputs["state"] = state
         # Everything else the caller passed goes through, minus what this model does
-        # not take. It used to be an allowlist of three names, which meant
-        # `generate(output_hidden_states=True)` returned a tuple of `None` -- the flag
-        # was dropped here, the forward never saw it, and generate collected the
+        # not take. An allowlist of specific names instead would make
+        # `generate(output_hidden_states=True)` return a tuple of `None`: the flag
+        # would be dropped here, the forward would never see it, and generate collects the
         # nothing it got back. Any user kwarg met the same fate, silently. The two
         # excluded here are `labels`, which would make generate compute a loss it
         # never reads, and the KV-cache bookkeeping that belongs to models with a KV

sparse_channel_mix.py

@@ -1,175 +0,0 @@
-# Copyright 2026 The RWKV team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Triton sparse channel-mix value projection for RWKV-7 decode.
-
-Kept out of `modeling_rwkv7.py` so the model file stays importable without Triton;
-`sparse_channel_mix_value` there dispatches here and falls back to a dense matmul.
-
-Registered as a custom op rather than called directly: a raw Triton launch is on
-dynamo's skip list, so calling it from the model breaks the graph at every layer
-(21 breaks in a 7.2B forward, measured) and costs about 8% end to end. As an opaque
-op it stays inside the single captured graph.
-
-The weight is indexed `[inter, hidden]` so one selected input channel reads one
-contiguous row.
-
-STRUCTURE. The surviving indices are compacted into a shared list first, and the
-projection then walks that list with scalar loads and a runtime trip count. Five
-structures were measured in situ on a 7.2B decode -- end to end, batch 1, under
-`max-autotune` with CUDA graphs engaged:
-
-  compacted indices, SPLIT=128              133.9 tok/s   <- shipped
-  tile-scan, (inter/16, hidden/1024) grid   127.8
-  scalar re-read instead of a reduction     slower than tile-scan
-  input axis moved off the grid, 512 prog   slower than tile-scan
-  coarse output tiles, few programs         far worse
-
-The compacted structure was tried once before and rejected, and why is worth
-keeping, because the reason was true and has stopped being true. Measured with one
-launch per stage and no CUDA graph, its extra launch and the barrier its
-data-dependent trip count forces cost ~1.3 ms/step of pipeline bubbles against the
-0.43 ms of kernel time it saved. Once the decode compiled into a single graph with
-cudagraphs engaged, GPU-busy time equals wall time and there are no bubbles left to
-pay -- so the arithmetic that killed it no longer applied to anything.
-
-The first re-test still lost, 122.3 against 127.8, because it inherited SPLIT=8:
-128 programs against the tile-scan's 4096. The split decides how much of the card
-the walk occupies and had never been swept. It is the whole difference:
-
-  SPLIT   8 -> 122.3      SPLIT  64 -> 133.3
-  SPLIT  16 -> 131.1      SPLIT 128 -> 133.9   <- optimum, bracketed both ways
-  SPLIT  32 -> 133.3      SPLIT 256 -> 132.9
-                          SPLIT 512 -> 132.4
-
-The durable lesson is about the record rather than the kernel: a rejection is only
-as good as the conditions it was measured under, and one that does not name them
-gets re-read as a fact about the algorithm. 128 is measured at hidden=4096; other
-widths have not been swept.
-
-Measure any variant of this end to end, not as a standalone kernel. Three separate
-standalone harnesses each favoured a configuration that lost in the model -- by
-allowing L2 reuse, by letting independent iterations overlap, and by using random
-weights whose `relu(x)**2` is ~50% dense instead of the ~10% a real checkpoint
-gives at any decode step.
-
-Cross-tile partials land in an fp32 accumulator through atomics, and the finalize
-pass re-zeros both that accumulator and the counter as it casts, so no separate
-clear is ever launched.
-"""
-
-import torch
-from transformers.utils.import_utils import is_triton_available
-
-# Guarded rather than bare, although the module itself is only ever imported
-# under the same `is_triton_available()` condition: import scanners that parse
-# this file in isolation (transformers' remote-code `check_imports` is one)
-# exempt imports inside an availability-guarded block, and demand triton on
-# every machine otherwise.
-if is_triton_available():
-    import triton
-    import triton.language as tl
-
-
-# How many programs the walk is split into. Measured optimum at hidden=4096; it sets
-# how much of the card the matvec occupies, so it is the first thing to sweep on a
-# model of a different width.
-_SPLIT = 128
-_BLOCK_H = 256
-_BLOCK_C = 256
-
-
-@triton.jit
-def _compact_kernel(act_ptr, idx_ptr, val_ptr, cnt_ptr, inter, BLOCK: tl.constexpr):
-    """Append this tile's surviving (index, value) pairs to one shared list."""
-    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
-    act = tl.load(act_ptr + offs, mask=offs < inter, other=0.0).to(tl.float32)
-    nonzero = act != 0.0
-    count = tl.sum(nonzero.to(tl.int32))
-    if count > 0:
-        # reserve a contiguous range, then place each survivor at its rank within it
-        base = tl.atomic_add(cnt_ptr, count)
-        rank = tl.cumsum(nonzero.to(tl.int32)) - nonzero.to(tl.int32)
-        tl.store(idx_ptr + base + rank, offs, mask=nonzero)
-        tl.store(val_ptr + base + rank, act, mask=nonzero)
-
-
-@triton.jit
-def _sparse_value_kernel(
-    idx_ptr,
-    val_ptr,
-    cnt_ptr,
-    w_ptr,
-    acc_ptr,
-    hidden,
-    SPLIT: tl.constexpr,
-    BLOCK_H: tl.constexpr,
-):
-    """Walk a slice of the compacted list, accumulating into one output tile."""
-    pid_h, pid_s = tl.program_id(0), tl.program_id(1)
-    total = tl.load(cnt_ptr)
-    per = (total + SPLIT - 1) // SPLIT
-    start = pid_s * per
-    stop = tl.minimum(start + per, total)
-    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
-    mask_h = offs_h < hidden
-    acc = tl.zeros([BLOCK_H], dtype=tl.float32)
-    for i in range(start, stop):
-        idx = tl.load(idx_ptr + i)
-        a = tl.load(val_ptr + i)
-        acc += a * tl.load(w_ptr + idx * hidden + offs_h, mask=mask_h, other=0.0).to(tl.float32)
-    if stop > start:
-        tl.atomic_add(acc_ptr + offs_h, acc, mask=mask_h)
-
-
-@triton.jit
-def _sparse_finalize_kernel(acc_ptr, out_ptr, cnt_ptr, hidden, BLOCK: tl.constexpr):
-    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
-    mask = offs < hidden
-    value = tl.load(acc_ptr + offs, mask=mask, other=0.0)
-    tl.store(out_ptr + offs, value.to(out_ptr.dtype.element_ty), mask=mask)
-    tl.store(acc_ptr + offs, tl.zeros([BLOCK], dtype=tl.float32), mask=mask)
-    if tl.program_id(0) == 0:
-        tl.store(cnt_ptr, 0)
-
-
-@torch.library.custom_op("rwkv7::sparse_channel_mix_value", mutates_args={"accumulator", "index", "value", "counter"})
-def triton_sparse_value(
-    activation: torch.Tensor,
-    weight_t: torch.Tensor,
-    accumulator: torch.Tensor,
-    index: torch.Tensor,
-    value: torch.Tensor,
-    counter: torch.Tensor,
-) -> torch.Tensor:
-    inter, hidden = weight_t.shape
-    out = torch.empty(hidden, device=activation.device, dtype=activation.dtype)
-    _compact_kernel[(triton.cdiv(inter, _BLOCK_C),)](activation, index, value, counter, inter, BLOCK=_BLOCK_C)
-    _sparse_value_kernel[(triton.cdiv(hidden, _BLOCK_H), _SPLIT)](
-        index,
-        value,
-        counter,
-        weight_t,
-        accumulator,
-        hidden,
-        SPLIT=_SPLIT,
-        BLOCK_H=_BLOCK_H,
-        num_warps=4,
-    )
-    _sparse_finalize_kernel[(triton.cdiv(hidden, 256),)](accumulator, out, counter, hidden, BLOCK=256)
-    return out
-
-
-@triton_sparse_value.register_fake
-def _(activation, weight_t, accumulator, index, value, counter):
-    return activation.new_empty(weight_t.shape[1])

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|rwkv_tokenizer_end_of_text|>",
+  "eos_token": "\n\n",
+  "unk_token": "<|rwkv_tokenizer_end_of_text|>",
+  "pad_token": "<|rwkv_tokenizer_end_of_text|>"
+}

tokenizer_config.json

@@ -0,0 +1,28 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|rwkv_tokenizer_end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "hf_rwkv_tokenizer.RwkvTokenizer",
+      null
+    ]
+  },
+  "bos_token": "<|rwkv_tokenizer_end_of_text|>",
+  "pad_token": "<|rwkv_tokenizer_end_of_text|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "\n\n",
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "RwkvTokenizer",
+  "unk_token": "<|rwkv_tokenizer_end_of_text|>",
+  "use_fast": false,
+  "chat_template": "{{ '<|rwkv_tokenizer_end_of_text|>' }}{% for message in messages %}{% if message['role'] == 'user' %}{{'User: ' + message['content'] + '\n\n'}}{% elif message['role'] == 'system' %}{{'System: ' + message['content'] + '\n\n'}}{% elif message['role'] == 'assistant' %}{{'Assistant: ' + message['content'] + '\n\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
+}