Upload folder using huggingface_hub

.ipynb_checkpoints/configuration_rwkv7-checkpoint.py

@@ -0,0 +1,65 @@
+########################################################################################################
+# RWKV-7 "Goose" (x070 / g1d) HuggingFace configuration
+# Based on the reference implementation from https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class RWKV7Config(PretrainedConfig):
+    """Configuration for the RWKV-7 (x070 / g1d) language model.
+
+    The defaults match the ``rwkv7-g1d-0.1b`` checkpoint
+    (L12-D768, head_size 64), but the embedding / lm-head vocabulary has been
+    re-sized (and re-initialized) to match the OLMo tokenizer.
+    """
+
+    model_type = "rwkv7"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=100278,
+        hidden_size=768,
+        num_hidden_layers=12,
+        head_size=64,
+        intermediate_size=3072,
+        decay_lora=64,
+        aaa_lora=64,
+        mv_lora=32,
+        gate_lora=128,
+        layer_norm_epsilon=1e-5,
+        group_norm_epsilon=64e-5,
+        bos_token_id=None,
+        eos_token_id=100257,
+        pad_token_id=100277,
+        tie_word_embeddings=False,
+        use_cuda_kernel=True,
+        chunk_len=16,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.head_size = head_size
+        self.intermediate_size = intermediate_size
+        self.decay_lora = decay_lora
+        self.aaa_lora = aaa_lora
+        self.mv_lora = mv_lora
+        self.gate_lora = gate_lora
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.group_norm_epsilon = group_norm_epsilon
+        # attention/ffn dims are derived from hidden_size in the reference model
+        self.attention_hidden_size = hidden_size
+        self.use_cuda_kernel = use_cuda_kernel
+        self.chunk_len = chunk_len
+
+        assert hidden_size % head_size == 0, "hidden_size must be divisible by head_size"
+
+        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,
+        )

README.md

@@ -0,0 +1,173 @@
+# RWKV-7 "g1d" 0.1B — HuggingFace port (OLMo tokenizer)
+
+A `trust_remote_code=True` HuggingFace wrapper around the BlinkDL
+**RWKV-7 "Goose" g1d 0.1B** checkpoint
+(`rwkv7-g1d-0.1b-20260129-ctx8192.pth`), re-headed for the **OLMo tokenizer**.
+
+- The 12 transformer-style RWKV-7 blocks keep their **pretrained** weights.
+- The **embedding** and **lm-head** are **re-initialized** (from RWKV's 65536-token
+  vocab to OLMo's 100278-token vocab), so they are *untrained* and need fine-tuning.
+- Time-mixing runs on a **fused CUDA kernel** (forward **and** backward) when a GPU
+  + CUDA toolchain are available, and transparently **falls back to pure PyTorch**
+  otherwise.
+
+> Because emb/head are freshly initialized, `generate()` produces gibberish until
+> you fine-tune on OLMo-tokenized data. The model *body* is pretrained; only the
+> vocabulary projection is new.
+
+---
+
+## Files
+
+| File | Purpose |
+|------|---------|
+| `config.json` | Serialized `RWKV7Config` (dims + `auto_map` to the remote-code classes). |
+| `configuration_rwkv7.py` | `RWKV7Config` — all architecture hyperparameters. |
+| `modeling_rwkv7.py` | Model code: CUDA-kernel dispatch + PyTorch fallback, `RWKV7Model`, `RWKV7ForCausalLM`. |
+| `cuda/wkv7_cuda.cu`, `cuda/wkv7_op.cpp` | The fused bf16 "wind_backstepping" RWKV-7 kernel (forward + backward), copied from `RWKV-v7/train_temp/cuda/`. |
+| `model.safetensors` | Converted weights (bf16, ~244M params). |
+| `tokenizer.json`, `tokenizer_config.json` | OLMo tokenizer (vocab 100278, GPT2-style BPE). |
+| `generation_config.json` | Default generation settings (eos/pad ids). |
+| `convert.py` | Reproduces `model.safetensors` from the original `.pth`. |
+| `verify.py` | End-to-end smoke test (load / forward / parity / backward / generate). |
+
+## Architecture (from `config.json`)
+
+| field | value | meaning |
+|-------|-------|---------|
+| `num_hidden_layers` | 12 | RWKV-7 blocks |
+| `hidden_size` | 768 | embedding dim `C` |
+| `head_size` | 64 | → 12 heads (`H = C / head_size`) |
+| `intermediate_size` | 3072 | channel-mix (FFN) hidden |
+| `decay_lora` / `aaa_lora` / `mv_lora` / `gate_lora` | 64 / 64 / 32 / 128 | LoRA ranks for `w` / `a` / `v` / `g` |
+| `vocab_size` | 100278 | OLMo tokenizer size (re-initialized emb/head) |
+| `chunk_len` | 16 | CUDA kernel chunk length; sequence is padded to a multiple of this |
+| `use_cuda_kernel` | true | prefer the fused kernel when possible |
+
+---
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+PATH = "/workspace/rwkv7-g1d-olmo"
+
+tok = AutoTokenizer.from_pretrained(PATH, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    PATH, trust_remote_code=True, dtype=torch.bfloat16
+).cuda().eval()
+
+ids = tok("The Eiffel tower is in the city of", return_tensors="pt").input_ids.cuda()
+with torch.no_grad():
+    logits = model(ids).logits          # (1, T, 100278)
+
+# training / backward
+model.train()
+out = model(ids, labels=ids)            # shifted cross-entropy loss
+out.loss.backward()                     # gradients flow through the CUDA kernel
+```
+
+### Forcing the PyTorch fallback
+
+Set the flag on the config (useful on CPU, non-bf16, or to debug the kernel):
+
+```python
+model.config.use_cuda_kernel = False    # every RWKV-7 op now runs in pure PyTorch
+```
+
+The kernel is also skipped automatically when: no CUDA device, input dtype is not
+bfloat16, or the kernel fails to compile — in all cases the fallback is used and a
+one-line notice is printed.
+
+---
+
+## How the CUDA kernel + fallback work (`modeling_rwkv7.py`)
+
+The time-mixing recurrence is dispatched by `run_rwkv7(r, w, k, v, a, b, config)`:
+
+1. **Kernel path** (`_rwkv7_cuda`) — used when `config.use_cuda_kernel` is set, the
+   tensors are CUDA + bf16, and `_try_load_cuda_kernel()` succeeds. On first call it
+   JIT-compiles `cuda/wkv7_op.cpp` + `cuda/wkv7_cuda.cu` via
+   `torch.utils.cpp_extension.load` (flags `-D_C_=head_size`, `-D_CHUNK_LEN_=chunk_len`)
+   and registers the `wind_backstepping` op. Compilation happens once per process and
+   is cached; any failure is caught and flips the model to the fallback.
+   - `_WindBackstepping` is a `torch.autograd.Function`:
+     - **forward** → `wind_backstepping.forward` (produces `y` plus the saved state
+       `s` and `sa` needed for backprop).
+     - **backward** → `wind_backstepping.backward` (returns gradients for all six
+       inputs `w, q, k, v, z, b`), so training works end-to-end on the kernel.
+   - The sequence length is padded to a multiple of `chunk_len` (16) and sliced back.
+
+2. **Fallback path** (`_rwkv7_pytorch`) — a plain sequential-over-time
+   implementation of the same recurrence
+   `state = state*exp(-exp(w)) + state·aᵀ·b + vᵀ·k`, `y = state·r`, in fp32.
+   It is fully differentiable through ordinary autograd (no custom backward needed).
+
+Both paths take the **raw (pre-exp) decay `w`** and apply `exp(-exp(w))` internally,
+so they are numerically interchangeable. `verify.py` confirms kernel-vs-fallback
+parity (identical bf16 logits and top-1 prediction).
+
+### Model structure
+
+`RWKV7ForCausalLM` → `.rwkv` (`RWKV7Model`) + `.head` (lm-head). `RWKV7Model` holds
+`emb`, `blocks[0..11]` (`RWKV7Block` = `ln1` + `att` time-mix, `ln2` + `ffn`
+channel-mix; block 0 also has `ln0`), and `ln_out`. State-dict keys mirror the
+original RWKV layout under the `rwkv.` prefix (e.g. `rwkv.blocks.0.att.receptance.weight`).
+Generation runs in **GPT mode** (no incremental KV cache — the full sequence is
+recomputed each step; correct but not the fastest).
+
+---
+
+## Reproducing the conversion (`convert.py`)
+
+```bash
+# 1) download the original checkpoint (already done in /workspace)
+wget https://huggingface.co/BlinkDL/rwkv7-g1/resolve/main/rwkv7-g1d-0.1b-20260129-ctx8192.pth \
+     -O /workspace/rwkv7-g1d-0.1b-20260129-ctx8192.pth
+
+# 2) convert -> /workspace/rwkv7-g1d-olmo
+python convert.py
+```
+
+What `convert.py` does:
+
+1. Loads the OLMo tokenizer to read the target vocab size (100278).
+2. Loads the `.pth` and infers all dims from the tensor shapes → builds `RWKV7Config`.
+3. Remaps the original RWKV keys to the HF module layout (`rwkv.` prefix), **dropping**
+   `emb.weight` / `head.weight`, and loads them with `strict=False`
+   (asserts there are **no** unexpected or unmatched keys besides emb/head).
+4. **Re-initializes** the embedding (`uniform(±1e-4)`) and head
+   (`orthogonal`, gain `0.5·√(vocab/hidden)`) for the new vocabulary — RWKV's own
+   init scheme.
+5. Saves weights (`safetensors`), config, tokenizer, and copies the remote-code files.
+
+## Verifying (`verify.py`)
+
+```bash
+python verify.py
+```
+
+Checks: load via `AutoModelForCausalLM(trust_remote_code=True)`, CUDA-kernel forward,
+kernel-vs-fallback parity, a backward pass (gradients on emb / attention / decay-LoRA),
+and a short greedy `generate()`.
+
+> Note: 399/402 parameters receive gradients — the 3 without are
+> `blocks.0.att.v0/v1/v2` (the value-residual params are unused in layer 0 by design).
+
+---
+
+## Requirements
+
+- PyTorch with CUDA (bf16-capable GPU; tested on RTX 3060 / CUDA 13.0) for the kernel;
+  CPU/other works via the fallback.
+- `transformers >= 5`, `safetensors`.
+- A working CUDA toolchain (`nvcc`) for first-call kernel JIT compilation; if absent,
+  the model still runs on the PyTorch fallback.
+
+## Credits
+
+RWKV-7 architecture and the original checkpoint/kernels by **BlinkDL** —
+<https://github.com/BlinkDL/RWKV-LM>. Kernel sources copied from
+`RWKV-v7/train_temp/cuda/`.

config.json

@@ -0,0 +1,30 @@
+{
+  "aaa_lora": 64,
+  "architectures": [
+    "RWKV7ForCausalLM"
+  ],
+  "attention_hidden_size": 768,
+  "auto_map": {
+    "AutoConfig": "configuration_rwkv7.RWKV7Config",
+    "AutoModelForCausalLM": "modeling_rwkv7.RWKV7ForCausalLM"
+  },
+  "bos_token_id": null,
+  "chunk_len": 16,
+  "decay_lora": 64,
+  "dtype": "bfloat16",
+  "eos_token_id": 100257,
+  "gate_lora": 128,
+  "group_norm_epsilon": 0.00064,
+  "head_size": 64,
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "rwkv7",
+  "mv_lora": 32,
+  "num_hidden_layers": 12,
+  "pad_token_id": 100277,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.1",
+  "use_cuda_kernel": true,
+  "vocab_size": 100278
+}

configuration_rwkv7.py

@@ -0,0 +1,65 @@
+########################################################################################################
+# RWKV-7 "Goose" (x070 / g1d) HuggingFace configuration
+# Based on the reference implementation from https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class RWKV7Config(PretrainedConfig):
+    """Configuration for the RWKV-7 (x070 / g1d) language model.
+
+    The defaults match the ``rwkv7-g1d-0.1b`` checkpoint
+    (L12-D768, head_size 64), but the embedding / lm-head vocabulary has been
+    re-sized (and re-initialized) to match the OLMo tokenizer.
+    """
+
+    model_type = "rwkv7"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=100278,
+        hidden_size=768,
+        num_hidden_layers=12,
+        head_size=64,
+        intermediate_size=3072,
+        decay_lora=64,
+        aaa_lora=64,
+        mv_lora=32,
+        gate_lora=128,
+        layer_norm_epsilon=1e-5,
+        group_norm_epsilon=64e-5,
+        bos_token_id=None,
+        eos_token_id=100257,
+        pad_token_id=100277,
+        tie_word_embeddings=False,
+        use_cuda_kernel=True,
+        chunk_len=16,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.head_size = head_size
+        self.intermediate_size = intermediate_size
+        self.decay_lora = decay_lora
+        self.aaa_lora = aaa_lora
+        self.mv_lora = mv_lora
+        self.gate_lora = gate_lora
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.group_norm_epsilon = group_norm_epsilon
+        # attention/ffn dims are derived from hidden_size in the reference model
+        self.attention_hidden_size = hidden_size
+        self.use_cuda_kernel = use_cuda_kernel
+        self.chunk_len = chunk_len
+
+        assert hidden_size % head_size == 0, "hidden_size must be divisible by head_size"
+
+        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,
+        )

convert.py

@@ -0,0 +1,107 @@
+"""Convert the BlinkDL rwkv7-g1d-0.1b .pth checkpoint into a HuggingFace
+`trust_remote_code` model directory, re-sizing (and re-initializing) the
+embedding + lm-head to the OLMo tokenizer vocabulary.
+"""
+
+import os
+import sys
+import math
+import shutil
+
+import torch
+import torch.nn as nn
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from configuration_rwkv7 import RWKV7Config
+from modeling_rwkv7 import RWKV7ForCausalLM
+
+PTH = "/workspace/rwkv7-g1d-0.1b-20260129-ctx8192.pth"
+OLMO = "/workspace/olmo"
+OUT = "/workspace/rwkv7-g1d-olmo"
+
+
+def main():
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(OLMO, trust_remote_code=True)
+    new_vocab = len(tok)
+    print(f"OLMo tokenizer vocab = {new_vocab}")
+
+    sd = torch.load(PTH, map_location="cpu")
+    old_vocab, n_embd = sd["emb.weight"].shape
+    print(f"checkpoint: vocab={old_vocab}, n_embd={n_embd}")
+
+    config = RWKV7Config(
+        vocab_size=new_vocab,
+        hidden_size=n_embd,
+        num_hidden_layers=12,
+        head_size=64,
+        intermediate_size=sd["blocks.0.ffn.key.weight"].shape[0],
+        decay_lora=sd["blocks.0.att.w1"].shape[1],
+        aaa_lora=sd["blocks.0.att.a1"].shape[1],
+        mv_lora=sd["blocks.0.att.v1"].shape[1],
+        gate_lora=sd["blocks.0.att.g1"].shape[1],
+        eos_token_id=tok.eos_token_id,
+        pad_token_id=tok.pad_token_id,
+        torch_dtype="bfloat16",
+        architectures=["RWKV7ForCausalLM"],
+        auto_map={
+            "AutoConfig": "configuration_rwkv7.RWKV7Config",
+            "AutoModelForCausalLM": "modeling_rwkv7.RWKV7ForCausalLM",
+        },
+    )
+
+    # remap RWKV native keys -> HF module layout (rwkv.* + head)
+    remap = {}
+    for k, v in sd.items():
+        if k == "emb.weight" or k == "head.weight":
+            continue  # re-initialized below
+        if k.startswith("blocks.") or k.startswith("ln_out."):
+            remap["rwkv." + k] = v
+        else:
+            remap["rwkv." + k] = v
+
+    model = RWKV7ForCausalLM(config).to(dtype=torch.bfloat16)
+
+    missing, unexpected = model.load_state_dict(remap, strict=False)
+    # emb/head are expected to be "missing" from `remap` (we drop them on purpose)
+    missing = [m for m in missing if not (m == "rwkv.emb.weight" or m == "head.weight")]
+    print("unexpected keys:", unexpected)
+    print("missing (besides emb/head):", missing)
+    assert not unexpected, f"unexpected keys present: {unexpected}"
+    assert not missing, f"unmatched keys: {missing}"
+
+    # Re-initialize embedding + head for the new (OLMo) vocabulary, using the
+    # RWKV reference init scheme.
+    with torch.no_grad():
+        emb = torch.empty(new_vocab, n_embd)
+        nn.init.uniform_(emb, a=-1e-4, b=1e-4)
+        model.rwkv.emb.weight.copy_(emb.to(torch.bfloat16))
+
+        head = torch.empty(new_vocab, n_embd)
+        scale = 0.5 * math.sqrt(new_vocab / n_embd) if new_vocab > n_embd else 0.5
+        nn.init.orthogonal_(head, gain=scale)
+        model.head.weight.copy_(head.to(torch.bfloat16))
+    print(f"re-initialized emb {tuple(model.rwkv.emb.weight.shape)} and head "
+          f"{tuple(model.head.weight.shape)} (head gain {scale:.4f})")
+
+    os.makedirs(OUT, exist_ok=True)
+    model.save_pretrained(OUT, safe_serialization=True)
+    config.save_pretrained(OUT)
+
+    # tokenizer + generation config
+    tok.save_pretrained(OUT)
+
+    # make sure the remote-code files sit alongside the weights
+    for fn in ["configuration_rwkv7.py", "modeling_rwkv7.py"]:
+        src = os.path.join(os.path.dirname(os.path.abspath(__file__)), fn)
+        dst = os.path.join(OUT, fn)
+        if os.path.abspath(src) != os.path.abspath(dst):
+            shutil.copy(src, dst)
+
+    print("saved to", OUT)
+    print("files:", sorted(os.listdir(OUT)))
+
+
+if __name__ == "__main__":
+    main()

cuda/wkv7_cuda.cu

@@ -0,0 +1,138 @@
+#include <cuda_bf16.h>
+#include <assert.h>
+
+using bf = __nv_bfloat16;
+__device__ inline float to_float(const bf & u) { return __bfloat162float(u); }
+__device__ inline bf to_bf(const float & u) { return __float2bfloat16_rn(u); }
+
+typedef bf * __restrict__ F_;
+
+__global__ void forward_kernel(int T, int H, F_ w_, F_ q_, F_ k_, F_ v_, F_ a_, F_ b_, bf* y_, float* s_, float* sa_) {
+    constexpr int C = _C_;
+    int bb = blockIdx.y, hh = blockIdx.x, i = threadIdx.x;
+
+    float state[C] = {0};
+    __shared__ float q[C], k[C], w[C], a[C], b[C];
+
+    for (int t = 0; t < T; t++) {
+        int ind = bb*T*H*C + t*H*C + hh * C + i;
+        __syncthreads();
+        q[i] = to_float(q_[ind]);
+        w[i] = __expf(-__expf(to_float(w_[ind])));
+        k[i] = to_float(k_[ind]);
+        a[i] = to_float(a_[ind]);
+        b[i] = to_float(b_[ind]);
+        __syncthreads();
+
+        float sa = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            sa += a[j] * state[j];
+        }
+        sa_[ind] = sa;
+
+        float v = to_float(v_[ind]);
+        float y = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            float& s = state[j];
+            s = s * w[j] + sa * b[j] + k[j] * v;
+            y += s * q[j];
+        }
+        y_[ind] = to_bf(y);
+
+        if ((t+1)%_CHUNK_LEN_ == 0) {
+            int base = (bb*H+hh)*(T/_CHUNK_LEN_)*C*C + (t/_CHUNK_LEN_)*C*C + i;
+#pragma unroll
+            for (int j = 0; j < C; j++) {
+                s_[base + j*C] = state[j];
+            }
+        }
+    }
+}
+
+__global__ void backward_kernel(int T, int H, F_ w_, F_ q_, F_ k_, F_ v_, F_ a_, F_ b_, F_ dy_, float * __restrict__ s_, float * __restrict__ sa_, bf* dw_, bf* dq_, bf* dk_, bf* dv_, bf* da_, bf* db_) {
+    constexpr int C = _C_;
+    int bb = blockIdx.y, hh = blockIdx.x, i = threadIdx.x;
+
+    float stateT[C] = {0}, dstate[C] = {0}, dstateT[C] = {0};
+    __shared__ float w[C], q[C], k[C], v[C], a[C], b[C], dy[C], sa[C], dSb_shared[C];
+    float qi, wi, ki, ai, bi, dyi;
+
+    for (int t = T-1; t >= 0; t--) {
+        int ind = bb*T*H*C + t*H*C + hh * C + i;
+        __syncthreads();
+        q[i] = qi = to_float(q_[ind]);
+        float wi_fac = -__expf(to_float(w_[ind]));
+        w[i] = wi = __expf(wi_fac);
+        k[i] = ki = to_float(k_[ind]);
+        a[i] = ai = to_float(a_[ind]);
+        b[i] = bi = to_float(b_[ind]);
+        v[i] = to_float(v_[ind]);
+        dy[i] = dyi = to_float(dy_[ind]);
+        sa[i] = sa_[ind];
+        __syncthreads();
+
+        if ((t+1)%_CHUNK_LEN_ == 0) {
+            int base = (bb*H+hh)*(T/_CHUNK_LEN_)*C*C + (t/_CHUNK_LEN_)*C*C + i*C;
+#pragma unroll
+            for (int j = 0; j < C; j++) {
+                stateT[j] = s_[base + j];
+            }
+        }
+
+        float dq = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            dq += stateT[j]*dy[j];
+        }
+        dq_[ind] = to_bf(dq);
+
+        float iwi = 1.0f/wi;
+#pragma unroll        
+        for (int j = 0; j < C; j++) {
+            stateT[j] = (stateT[j] - ki*v[j] - bi*sa[j]) * iwi;
+            dstate[j] += dyi * q[j];
+            dstateT[j] += qi * dy[j];
+        }
+
+        float dw = 0, dk = 0, dv = 0, db = 0, dSb = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            dw += dstateT[j]*stateT[j];
+            dk += dstateT[j]*v[j];
+            dv += dstate[j]*k[j];
+            dSb += dstate[j]*b[j];
+            db += dstateT[j]*sa[j];
+        }
+        dw_[ind] = to_bf(dw * wi * wi_fac);
+        dk_[ind] = to_bf(dk);
+        dv_[ind] = to_bf(dv);
+        db_[ind] = to_bf(db);
+
+        __syncthreads();
+        dSb_shared[i] = dSb;
+        __syncthreads();
+
+        float da = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            da += stateT[j]*dSb_shared[j];
+        }
+        da_[ind] = to_bf(da);
+
+#pragma unroll        
+        for (int j = 0; j < C; j++) {
+            dstate[j] = dstate[j]*w[j] + dSb * a[j];
+            dstateT[j] = dstateT[j]*wi + ai * dSb_shared[j];
+        }
+    }
+}
+
+void cuda_forward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*y, float*s, float*sa) {
+    forward_kernel<<<dim3(H,B), dim3(_C_)>>>(T,H,w,q,k,v,z,a,y,s,sa);
+}
+void cuda_backward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*dy, float*s, float*sa, bf*dw, bf*dq, bf*dk, bf*dv, bf*dz, bf*da) {
+    assert(T%_CHUNK_LEN_ == 0);
+    backward_kernel<<<dim3(H,B), dim3(_C_)>>>(T,H,w,q,k,v,z,a,dy,s,sa,dw,dq,dk,dv,dz,da);
+}

cuda/wkv7_op.cpp

@@ -0,0 +1,29 @@
+#include <torch/extension.h>
+#include <cuda_bf16.h>
+using bf = __nv_bfloat16;
+
+void cuda_forward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*y, float*s, float*sa);
+
+void forward(torch::Tensor &w, torch::Tensor &q, torch::Tensor &k, torch::Tensor &v, torch::Tensor &z, torch::Tensor &a, torch::Tensor &y, torch::Tensor &s, torch::Tensor &sa) {
+    int B = w.sizes()[0], T = w.sizes()[1], H = w.sizes()[2];
+    cuda_forward(B, T, H, (bf*)w.data_ptr(), (bf*)q.data_ptr(), (bf*)k.data_ptr(), (bf*)v.data_ptr(), (bf*)z.data_ptr(), (bf*)a.data_ptr(), (bf*)y.data_ptr(), (float*)s.data_ptr(), (float*)sa.data_ptr());
+}
+
+void cuda_backward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*dy, float*s, float*sa, bf*dw, bf*dq, bf*dk, bf*dv, bf*dz, bf*da);
+
+void backward(torch::Tensor &w, torch::Tensor &q, torch::Tensor &k, torch::Tensor &v, torch::Tensor &z, torch::Tensor &a, torch::Tensor &dy,
+        torch::Tensor &s, torch::Tensor &sa, torch::Tensor &dw, torch::Tensor &dq, torch::Tensor &dk, torch::Tensor &dv, torch::Tensor &dz, torch::Tensor &da) {
+    int B = w.sizes()[0], T = w.sizes()[1], H = w.sizes()[2];
+    cuda_backward(B, T, H, (bf*)w.data_ptr(), (bf*)q.data_ptr(), (bf*)k.data_ptr(), (bf*)v.data_ptr(), (bf*)z.data_ptr(), (bf*)a.data_ptr(), (bf*)dy.data_ptr(), 
+            (float*)s.data_ptr(), (float*)sa.data_ptr(), (bf*)dw.data_ptr(), (bf*)dq.data_ptr(), (bf*)dk.data_ptr(), (bf*)dv.data_ptr(), (bf*)dz.data_ptr(), (bf*)da.data_ptr());
+}
+
+TORCH_LIBRARY(wind_backstepping, m) {
+    m.def("forward(Tensor w, Tensor q, Tensor k, Tensor v, Tensor z, Tensor a, Tensor(a!) y, Tensor(b!) s, Tensor(c!) sa) -> ()");
+    m.def("backward(Tensor w, Tensor q, Tensor k, Tensor v, Tensor z, Tensor a, Tensor dy, Tensor s, Tensor sa, Tensor(a!) dw, Tensor(b!) dq, Tensor(c!) dk, Tensor(d!) dv, Tensor(e!) dz, Tensor(f!) da) -> ()");
+}
+
+TORCH_LIBRARY_IMPL(wind_backstepping, CUDA, m) {
+    m.impl("forward", &forward);
+    m.impl("backward", &backward);
+}

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 100257,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 100277,
+  "transformers_version": "5.14.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a56cd7c04a4265fda944e60f27c63f2b7875fe6b1d3cb636f63ade926eb44d1d
+size 488935616

modeling_rwkv7.py

@@ -0,0 +1,464 @@
+########################################################################################################
+# RWKV-7 "Goose" (x070 / g1d) HuggingFace modeling code
+# Based on the reference implementation from https://github.com/BlinkDL/RWKV-LM
+#
+# This file provides a `trust_remote_code=True` compatible RWKV-7 model that:
+#   * uses a fused CUDA kernel (fwd + bwd, "wind_backstepping" bf16) when a GPU
+#     + a working CUDA toolchain are available, and
+#   * transparently falls back to a pure-PyTorch implementation otherwise
+#     (autograd handles the backward pass automatically in the fallback).
+########################################################################################################
+
+import os
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.generation import GenerationMixin
+
+try:
+    from .configuration_rwkv7 import RWKV7Config
+except ImportError:  # allow running as a plain script (e.g. conversion)
+    from configuration_rwkv7 import RWKV7Config
+
+
+########################################################################################################
+# CUDA kernel loading (lazy, best-effort).  If anything goes wrong we silently
+# fall back to the pure-PyTorch path.
+########################################################################################################
+
+_KERNEL_STATE = {"loaded": False, "ok": False, "op": None, "chunk_len": 16, "head_size": 64}
+
+
+def _try_load_cuda_kernel(head_size: int, chunk_len: int):
+    """Compile & register the wind_backstepping RWKV-7 CUDA op. Returns True on success."""
+    if _KERNEL_STATE["loaded"]:
+        return _KERNEL_STATE["ok"]
+    _KERNEL_STATE["loaded"] = True
+    _KERNEL_STATE["head_size"] = head_size
+    _KERNEL_STATE["chunk_len"] = chunk_len
+
+    if not torch.cuda.is_available():
+        _KERNEL_STATE["ok"] = False
+        return False
+
+    try:
+        from torch.utils.cpp_extension import load
+
+        this_dir = os.path.dirname(os.path.abspath(__file__))
+        cuda_dir = os.path.join(this_dir, "cuda")
+        sources = [
+            os.path.join(cuda_dir, "wkv7_op.cpp"),
+            os.path.join(cuda_dir, "wkv7_cuda.cu"),
+        ]
+        if not all(os.path.exists(s) for s in sources):
+            _KERNEL_STATE["ok"] = False
+            return False
+
+        flags = [
+            "-res-usage",
+            f"-D_C_={head_size}",
+            f"-D_CHUNK_LEN_={chunk_len}",
+            "--use_fast_math",
+            "-O3",
+            "-Xptxas -O3",
+            "--extra-device-vectorization",
+        ]
+        load(
+            name=f"wind_backstepping_c{head_size}_l{chunk_len}",
+            sources=sources,
+            is_python_module=False,
+            verbose=False,
+            extra_cuda_cflags=flags,
+        )
+        _KERNEL_STATE["op"] = torch.ops.wind_backstepping
+        _KERNEL_STATE["ok"] = True
+        return True
+    except Exception as e:  # noqa: BLE001 - any failure -> fallback
+        print(f"[RWKV7] CUDA kernel unavailable, using PyTorch fallback ({type(e).__name__}: {e})")
+        _KERNEL_STATE["ok"] = False
+        return False
+
+
+class _WindBackstepping(torch.autograd.Function):
+    """Fused RWKV-7 kernel wrapper (bf16). Implements both forward and backward.
+
+    Inputs are shaped (B, T, H, C) with T % CHUNK_LEN == 0 and dtype bfloat16.
+    """
+
+    @staticmethod
+    def forward(ctx, w, q, k, v, z, b):
+        op = _KERNEL_STATE["op"]
+        chunk_len = _KERNEL_STATE["chunk_len"]
+        B, T, H, C = w.shape
+        assert T % chunk_len == 0, "pad T to a multiple of CHUNK_LEN"
+        assert all(i.dtype == torch.bfloat16 for i in [w, q, k, v, z, b])
+        assert all(i.is_contiguous() for i in [w, q, k, v, z, b])
+        y = torch.empty_like(v)
+        s = torch.empty(B, H, T // chunk_len, C, C, dtype=torch.float32, device=w.device)
+        sa = torch.empty(B, T, H, C, dtype=torch.float32, device=w.device)
+        op.forward(w, q, k, v, z, b, y, s, sa)
+        ctx.save_for_backward(w, q, k, v, z, b, s, sa)
+        return y
+
+    @staticmethod
+    def backward(ctx, dy):
+        op = _KERNEL_STATE["op"]
+        assert dy.dtype == torch.bfloat16
+        dy = dy.contiguous()
+        w, q, k, v, z, b, s, sa = ctx.saved_tensors
+        dw, dq, dk, dv, dz, db = [torch.empty_like(x) for x in [w, q, k, v, z, b]]
+        op.backward(w, q, k, v, z, b, dy, s, sa, dw, dq, dk, dv, dz, db)
+        return dw, dq, dk, dv, dz, db
+
+
+def _rwkv7_cuda(r, w, k, v, a, b, head_size, chunk_len):
+    """CUDA path. r,w,k,v,a,b are (B, T, C) bf16.  a = -kk, b = kk*a_gate."""
+    B, T, C = r.shape
+    H = C // head_size
+    pad = (chunk_len - T % chunk_len) % chunk_len
+    if pad:
+        r, w, k, v, a, b = [F.pad(x, (0, 0, 0, pad)) for x in (r, w, k, v, a, b)]
+    Tp = T + pad
+    r, w, k, v, a, b = [x.view(B, Tp, H, head_size).contiguous() for x in (r, w, k, v, a, b)]
+    y = _WindBackstepping.apply(w, r, k, v, a, b).view(B, Tp, C)
+    if pad:
+        y = y[:, :T]
+    return y
+
+
+def _rwkv7_pytorch(r, w, k, v, a, b, head_size):
+    """Pure-PyTorch reference (sequential over time). Differentiable via autograd.
+
+    w is the raw (pre-exp) log-decay; the recurrence uses exp(-exp(w)).
+    """
+    B, T, C = r.size()
+    H = C // head_size
+    N = head_size
+    dtype_in = r.dtype
+    r = r.view(B, T, H, N).float()
+    k = k.view(B, T, H, N).float()
+    v = v.view(B, T, H, N).float()
+    a = a.view(B, T, H, N).float()
+    b = b.view(B, T, H, N).float()
+    w = torch.exp(-torch.exp(w.view(B, T, H, N).float()))
+
+    out = torch.zeros((B, T, H, N), device=r.device, dtype=torch.float32)
+    state = torch.zeros((B, H, N, N), device=r.device, dtype=torch.float32)
+    for t in range(T):
+        kk = k[:, t, :].view(B, H, 1, N)
+        rr = r[:, t, :].view(B, H, N, 1)
+        vv = v[:, t, :].view(B, H, N, 1)
+        aa = a[:, t, :].view(B, H, N, 1)
+        bb = b[:, t, :].view(B, H, 1, N)
+        state = state * w[:, t, :, None, :] + state @ aa @ bb + vv @ kk
+        out[:, t, :] = (state @ rr).view(B, H, N)
+    return out.view(B, T, C).to(dtype=dtype_in)
+
+
+def run_rwkv7(r, w, k, v, a, b, config, force_fallback=False):
+    """Dispatch to CUDA kernel when possible, otherwise PyTorch fallback."""
+    use_kernel = (
+        config.use_cuda_kernel
+        and not force_fallback
+        and r.is_cuda
+        and r.dtype == torch.bfloat16
+        and _try_load_cuda_kernel(config.head_size, config.chunk_len)
+    )
+    if use_kernel:
+        return _rwkv7_cuda(r, w, k, v, a, b, config.head_size, config.chunk_len)
+    return _rwkv7_pytorch(r, w, k, v, a, b, config.head_size)
+
+
+########################################################################################################
+# RWKV-7 time-mixing ("attention") block
+########################################################################################################
+
+
+class RWKV7TimeMix(nn.Module):
+    def __init__(self, config: RWKV7Config, layer_id: int):
+        super().__init__()
+        self.config = config
+        self.layer_id = layer_id
+        self.head_size = config.head_size
+        C = config.hidden_size
+        self.n_head = C // self.head_size
+        H, N = self.n_head, self.head_size
+
+        self.x_r = nn.Parameter(torch.empty(1, 1, C))
+        self.x_w = nn.Parameter(torch.empty(1, 1, C))
+        self.x_k = nn.Parameter(torch.empty(1, 1, C))
+        self.x_v = nn.Parameter(torch.empty(1, 1, C))
+        self.x_a = nn.Parameter(torch.empty(1, 1, C))
+        self.x_g = nn.Parameter(torch.empty(1, 1, C))
+
+        self.w0 = nn.Parameter(torch.empty(1, 1, C))
+        self.w1 = nn.Parameter(torch.empty(C, config.decay_lora))
+        self.w2 = nn.Parameter(torch.empty(config.decay_lora, C))
+
+        self.a0 = nn.Parameter(torch.empty(1, 1, C))
+        self.a1 = nn.Parameter(torch.empty(C, config.aaa_lora))
+        self.a2 = nn.Parameter(torch.empty(config.aaa_lora, C))
+
+        self.v0 = nn.Parameter(torch.empty(1, 1, C))
+        self.v1 = nn.Parameter(torch.empty(C, config.mv_lora))
+        self.v2 = nn.Parameter(torch.empty(config.mv_lora, C))
+
+        self.g1 = nn.Parameter(torch.empty(C, config.gate_lora))
+        self.g2 = nn.Parameter(torch.empty(config.gate_lora, C))
+
+        self.k_k = nn.Parameter(torch.empty(1, 1, C))
+        self.k_a = nn.Parameter(torch.empty(1, 1, C))
+        self.r_k = nn.Parameter(torch.empty(H, N))
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        self.receptance = nn.Linear(C, C, bias=False)
+        self.key = nn.Linear(C, C, bias=False)
+        self.value = nn.Linear(C, C, bias=False)
+        self.output = nn.Linear(C, C, bias=False)
+        self.ln_x = nn.GroupNorm(H, C, eps=config.group_norm_epsilon)
+
+    def forward(self, x, v_first):
+        B, T, C = x.size()
+        H = self.n_head
+        xx = self.time_shift(x) - x
+
+        xr = x + xx * self.x_r
+        xw = x + xx * self.x_w
+        xk = x + xx * self.x_k
+        xv = x + xx * self.x_v
+        xa = x + xx * self.x_a
+        xg = x + xx * self.x_g
+
+        r = self.receptance(xr)
+        # soft-clamp to (-inf, -0.5); the recurrence applies exp(-exp(w))
+        w = -F.softplus(-(self.w0 + torch.tanh(xw @ self.w1) @ self.w2)) - 0.5
+        k = self.key(xk)
+        v = self.value(xv)
+        if self.layer_id == 0:
+            v_first = v
+        else:
+            v = v + (v_first - v) * torch.sigmoid(self.v0 + (xv @ self.v1) @ self.v2)
+        a = torch.sigmoid(self.a0 + (xa @ self.a1) @ self.a2)  # in-context learning rate
+        g = torch.sigmoid(xg @ self.g1) @ self.g2
+
+        kk = k * self.k_k
+        kk = F.normalize(kk.view(B, T, H, -1), dim=-1, p=2.0).view(B, T, C)
+        k = k * (1 + (a - 1) * self.k_a)
+
+        x = run_rwkv7(r, w, k, v, -kk, kk * a, self.config)
+        x = self.ln_x(x.view(B * T, C)).view(B, T, C)
+
+        x = x + (
+            (r.view(B, T, H, -1) * k.view(B, T, H, -1) * self.r_k).sum(dim=-1, keepdim=True)
+            * v.view(B, T, H, -1)
+        ).view(B, T, C)
+        x = self.output(x * g)
+        return x, v_first
+
+
+########################################################################################################
+# RWKV-7 channel-mixing (FFN) block
+########################################################################################################
+
+
+class RWKV7ChannelMix(nn.Module):
+    def __init__(self, config: RWKV7Config, layer_id: int):
+        super().__init__()
+        self.layer_id = layer_id
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        self.x_k = nn.Parameter(torch.empty(1, 1, config.hidden_size))
+        self.key = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.value = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        xx = self.time_shift(x) - x
+        k = x + xx * self.x_k
+        k = torch.relu(self.key(k)) ** 2
+        return self.value(k)
+
+
+class RWKV7Block(nn.Module):
+    def __init__(self, config: RWKV7Config, layer_id: int):
+        super().__init__()
+        self.layer_id = layer_id
+        eps = config.layer_norm_epsilon
+        if layer_id == 0:
+            self.ln0 = nn.LayerNorm(config.hidden_size, eps=eps)
+        self.ln1 = nn.LayerNorm(config.hidden_size, eps=eps)
+        self.ln2 = nn.LayerNorm(config.hidden_size, eps=eps)
+        self.att = RWKV7TimeMix(config, layer_id)
+        self.ffn = RWKV7ChannelMix(config, layer_id)
+
+    def forward(self, x, v_first):
+        if self.layer_id == 0:
+            x = self.ln0(x)
+        x_attn, v_first = self.att(self.ln1(x), v_first)
+        x = x + x_attn
+        x = x + self.ffn(self.ln2(x))
+        return x, v_first
+
+
+########################################################################################################
+# HuggingFace wrappers
+########################################################################################################
+
+
+class RWKV7PreTrainedModel(PreTrainedModel):
+    config_class = RWKV7Config
+    base_model_prefix = "rwkv"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["RWKV7Block"]
+
+    def _init_weights(self, module):
+        # Weights normally come from a pretrained checkpoint; this only covers
+        # freshly-created (e.g. re-sized embedding / head) parameters.
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=1e-4)
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+
+class RWKV7Model(RWKV7PreTrainedModel):
+    def __init__(self, config: RWKV7Config):
+        super().__init__(config)
+        self.config = config
+        self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.blocks = nn.ModuleList(
+            [RWKV7Block(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.ln_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.emb
+
+    def set_input_embeddings(self, value):
+        self.emb = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else True
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.emb(input_ids)
+        x = inputs_embeds
+
+        all_hidden_states = () if output_hidden_states else None
+        v_first = torch.empty_like(x)
+        for block in self.blocks:
+            if output_hidden_states:
+                all_hidden_states += (x,)
+            if self.gradient_checkpointing and self.training:
+                x, v_first = self._gradient_checkpointing_func(
+                    block.__call__, x, v_first
+                )
+            else:
+                x, v_first = block(x, v_first)
+
+        x = self.ln_out(x)
+        if output_hidden_states:
+            all_hidden_states += (x,)
+
+        if not return_dict:
+            return tuple(v for v in [x, all_hidden_states] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=x,
+            hidden_states=all_hidden_states,
+        )
+
+
+class RWKV7ForCausalLM(RWKV7PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = []
+
+    def __init__(self, config: RWKV7Config):
+        super().__init__(config)
+        self.rwkv = RWKV7Model(config)
+        self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.rwkv.emb
+
+    def set_input_embeddings(self, value):
+        self.rwkv.emb = value
+
+    def get_output_embeddings(self):
+        return self.head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.head = new_embeddings
+
+    def get_decoder(self):
+        return self.rwkv
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else True
+
+        outputs = self.rwkv(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        hidden = outputs.last_hidden_state
+        logits = self.head(hidden)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)).float(),
+                shift_labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (logits,) + (outputs.hidden_states,) if outputs.hidden_states else (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, inputs_embeds=None, **kwargs):
+        # RWKV-7 here runs in GPT mode (no incremental KV cache); recompute the
+        # full sequence each step. Correct, though not the fastest.
+        model_inputs = {"input_ids": input_ids}
+        return model_inputs

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 65536,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|endoftext|>"
+}

verify.py

@@ -0,0 +1,68 @@
+"""End-to-end verification of the converted RWKV-7 HF model."""
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+PATH = "/workspace/rwkv7-g1d-olmo"
+
+print("=" * 70)
+print("1) Loading with AutoModelForCausalLM(trust_remote_code=True)")
+tok = AutoTokenizer.from_pretrained(PATH, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    PATH, trust_remote_code=True, dtype=torch.bfloat16
+).cuda()
+model.eval()
+print("   model class:", type(model).__name__)
+print("   emb:", tuple(model.get_input_embeddings().weight.shape),
+      "head:", tuple(model.get_output_embeddings().weight.shape))
+nparams = sum(p.numel() for p in model.parameters())
+print(f"   params: {nparams/1e6:.1f}M")
+
+ids = tok("The Eiffel tower is in the city of", return_tensors="pt").input_ids.cuda()
+print("   input ids:", ids.tolist())
+
+print("=" * 70)
+print("2) Forward with CUDA kernel")
+with torch.no_grad():
+    out_k = model(ids).logits
+print("   logits:", tuple(out_k.shape), out_k.dtype)
+
+print("=" * 70)
+print("3) Forward with PyTorch fallback + kernel-vs-fallback parity")
+model.config.use_cuda_kernel = False
+with torch.no_grad():
+    out_f = model(ids).logits
+model.config.use_cuda_kernel = True
+
+diff = (out_k.float() - out_f.float()).abs()
+print(f"   max abs diff kernel vs fallback: {diff.max().item():.4e}")
+print(f"   mean abs diff:                   {diff.mean().item():.4e}")
+# argmax agreement on last token
+print("   kernel top-1 next id:", out_k[0, -1].argmax().item(),
+      "| fallback:", out_f[0, -1].argmax().item())
+
+print("=" * 70)
+print("4) Backward pass (gradients flow through kernel)")
+model.train()
+ids2 = tok("Backward test sentence for gradient check.", return_tensors="pt").input_ids.cuda()
+out = model(ids2, labels=ids2)
+loss = out.loss
+loss.backward()
+gnorm_emb = model.get_input_embeddings().weight.grad
+g_att = model.rwkv.blocks[0].att.receptance.weight.grad
+g_w1 = model.rwkv.blocks[6].att.w1.grad
+print(f"   loss: {loss.item():.4f}")
+print(f"   emb.grad is not None: {gnorm_emb is not None}, norm={gnorm_emb.float().norm().item():.4e}")
+print(f"   blocks.0.att.receptance.grad norm: {g_att.float().norm().item():.4e}")
+print(f"   blocks.6.att.w1 (decay-lora) grad norm: {g_w1.float().norm().item():.4e}")
+n_with_grad = sum(1 for p in model.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
+n_total = sum(1 for _ in model.parameters())
+print(f"   params with non-zero grad: {n_with_grad}/{n_total}")
+
+print("=" * 70)
+print("5) generate()")
+model.eval()
+with torch.no_grad():
+    gen = model.generate(ids, max_new_tokens=10, do_sample=False)
+print("   generated:", tok.decode(gen[0]))
+print("=" * 70)
+print("ALL CHECKS DONE")