Upload build/torch-universal/rwkv7_kernel/__init__.py with huggingface_hub

build/torch-universal/rwkv7_kernel/__init__.py

@@ -0,0 +1,527 @@
+# 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.
+"""RWKV-7 WKV kernel, Hub layout.
+
+Complete-semantics entry point `rwkv7_wkv`: the single-token decode step runs a
+Triton kernel that reads the recurrent state once per (batch, head) instead of
+three times; every other shape -- prefill, packed rows, CPU, non-power-of-two
+head widths -- takes the portable reference paths, which are carried verbatim
+from the `transformers` RWKV-7 implementation (same author, same license) so
+this package is self-contained and correct on any input.
+
+Honest scope: the decode kernel's gain grows with batch size and is nil at
+batch 1 -- at batch 1 the step is not bound by state reads. Prefill is the
+reference chunk-parallel form, unaccelerated here.
+"""
+
+import math
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+
+    _HAVE_TRITON = True
+except ImportError:  # keep the module importable; the entry point falls back
+    _HAVE_TRITON = False
+
+
+_INV_SQRT_E = 0.6065306597126334
+
+
+# ---- portable reference paths (verbatim from transformers) ----
+
+
+def rwkv7_recurrent(
+    r: torch.Tensor,
+    w_log: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kk: torch.Tensor,
+    a: torch.Tensor,
+    state: torch.Tensor,
+    compute_dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Reference RWKV-7 WKV recurrence (generalised delta rule).
+
+    All inputs are `[batch, seq_len, num_heads, head_dim]`; `state` is
+    `[batch, num_heads, head_dim, head_dim]` and is carried across calls.
+
+    Per token, with `S` the per-head state and every product outer over the two
+    head_dim axes::
+
+        decay = exp(w_log)
+        sa    = (-kk) @ S                 # uses the PRE-update state
+        S     = decay * S + (kk * a) ⊗ sa + k ⊗ v
+        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:
+    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.
+    """
+    batch, seq_len, num_heads, head_dim = r.shape
+    dtype = r.dtype
+    r, w_log, k, v, kk, a = (t.to(compute_dtype) for t in (r, w_log, k, v, kk, a))
+    state = state.to(compute_dtype)
+    out = torch.empty(
+        batch, seq_len, num_heads, head_dim, device=r.device, dtype=compute_dtype
+    )
+
+    for t in range(seq_len):
+        decay = torch.exp(w_log[:, t])[..., None]  # [B, H, K, 1]
+        kk_t = kk[:, t]
+        b = (kk_t * a[:, t])[..., None]  # [B, H, K, 1]
+        # sa[b, h, j] = sum_i (-kk[b, h, i]) * S[b, h, i, j]
+        sa = torch.einsum("bhi,bhij->bhj", -kk_t, state)[:, :, None, :]  # [B, H, 1, V]
+        state = decay * state + b * sa + k[:, t][..., None] * v[:, t][:, :, None, :]
+        out[:, t] = torch.einsum("bhi,bhij->bhj", r[:, t], state)
+
+    return out.to(dtype), state
+
+
+def _unit_lower_triangular_inverse(
+    strict: torch.Tensor, block: int = 8
+) -> torch.Tensor:
+    """Inverse of a batch of unit lower triangular matrices `I + strict`.
+
+    `strict` is strictly lower triangular, so it is nilpotent and the Neumann series
+    `I - strict + strict^2 - ...` terminates -- which makes an inverse by repeated
+    squaring look like the obvious choice, and it is wrong here. The intermediate
+    powers are not bounded by the answer: on a real chunk from a 3-layer model at
+    T=1024, `strict` has entries at most 0.977 and the true inverse has entries at
+    most 1.0, but `strict^32` reaches 1.3e11. Summing a series whose terms exceed the
+    result by eleven orders of magnitude cancels away every digit float32 has -- both
+    Newton doubling and a plain Neumann sum came back with entries of 1e4 where the
+    answer is 1, and the NaNs that followed reached the states two layers on. It shows
+    up in a fifth of randomly initialised 3-layer models at T=1024, one in forty at
+    T=256 and none at all at T=128 or below: it needs several chunks to compound, so a
+    short-prompt test cannot see it however many initialisations it tries.
+
+    Block forward substitution never forms a high power. `block` bounds the largest
+    power taken, and the same chunk that breaks the series gives 6.6e-7 at block 8 and
+    1.3e-7 at block 4, against float64 `linalg.solve_triangular`; 16 already costs a
+    digit and a half (1.9e-5) because the within-block series runs longer.
+
+    Random triangular matrices do not show any of this, which is worth stating because
+    it is what a test would reach for first: they agree with the series to 1e-6, since
+    their own inverses are as large as the intermediate powers. The cancellation needs
+    an inverse that stays near 1 while the powers do not, which is what the delta rule
+    produces and what a random matrix does not.
+
+    `linalg.solve_triangular` is the obvious call and is the accurate one, but it
+    decomposes for export into a graph carrying a scalar tensor constant, which
+    `lift_fresh` turns into `aten.alias`, functionalization turns into the in-place
+    `aten.detach_`, and aot_autograd rejects -- one of the two ops that made every ONNX
+    export subtest fail. Inverting once for every chunk at a time, rather than solving
+    inside the serial chunk loop, is what keeps that affordable; it is not free even
+    so. Measured against the per-chunk solve, loop included, on an RTX 5090 at 1.5B
+    shapes: 2.5x at T=1024 where 16 chunks cannot fill the card, falling to 1.2x at
+    T=4096, and flat in batch and head count. End to end on CPU, a T=1024 forward
+    runs about a third slower than it did on the solve.
+    """
+    span = strict.shape[-1]
+    block = min(block, span)
+    # Padding with zeros extends the matrix by an identity block, whose inverse is
+    # itself, so the leading span x span block of the result is unchanged.
+    pad = -span % block
+    if pad:
+        strict = torch.nn.functional.pad(strict, (0, pad, 0, pad))
+    width = span + pad
+    blocks = width // block
+
+    eye = torch.eye(block, device=strict.device, dtype=strict.dtype)
+    # Every diagonal block at once: [..., blocks, block, block].
+    diag = strict.reshape(*strict.shape[:-2], blocks, block, blocks, block)
+    diag = diag.diagonal(dim1=-4, dim2=-2).movedim(-1, -3)
+    # Exact after `block` terms, and `block` is small enough that the powers stay near
+    # the answer -- which is the whole point of the blocking.
+    inverse, term = eye - diag, -diag
+    for _ in range(block - 2):
+        term = -(diag @ term)
+        inverse = inverse + term
+
+    # `out` grows one block row at a time and is the inverse of the leading principal
+    # submatrix throughout, so the substitution reads exactly the columns it has
+    # already filled and never multiplies the zeros above the diagonal.
+    out = inverse[..., 0, :, :]
+    for i in range(1, blocks):
+        start, stop = i * block, (i + 1) * block
+        d = inverse[..., i, :, :]
+        row = torch.cat([-d @ (strict[..., start:stop, :start] @ out), d], dim=-1)
+        out = torch.cat([torch.nn.functional.pad(out, (0, block)), row], dim=-2)
+    return out[..., :span, :span] if pad else out
+
+
+def rwkv7_chunked(
+    r: torch.Tensor,
+    w_log: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kk: torch.Tensor,
+    a: torch.Tensor,
+    state: torch.Tensor,
+    chunk_size: int = 64,
+    compute_dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Chunk-parallel form of [`rwkv7_recurrent`], for sequences.
+
+    The step is `S_t = A_t S_{t-1} + k_t v_t^T` with `A_t = diag(w_t) - b_t kk_t^T`
+    and `b_t = kk_t * a_t`, i.e. diagonal-plus-rank-one. Substituting
+    `S_t = diag(c_t) P_t` with the running decay `c_t = prod_{s<=t} w_s` removes the
+    diagonal part and leaves a plain delta rule::
+
+        P_t = (I - b~_t q~_t^T) P_{t-1} + k~_t v_t^T
+        b~ = b / c,  k~ = k / c,  q~ = kk * c_{t-1},  r~ = r * c
+
+    Writing `u_t = q~_t^T P_{t-1}` turns a whole chunk into one unit-lower-triangular
+    system, so the chunk is a handful of matmuls instead of `chunk_size` sequential
+    steps; only the chunk-to-chunk carry stays serial::
+
+        (I + tril(Q~ B~^T, -1)) U = Q~ P_0 + tril(Q~ K~^T, -1) V
+        O   = R~ P_0 + tril(R~ K~^T, 0) V - tril(R~ B~^T, 0) U
+        P_C = P_0 + K~^T V - B~^T U
+
+    `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
+    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
+    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
+    than the number led into the overflow band. `chunk_size` is checked below rather
+    than left to that reasoning.
+
+    Precision does not degrade with it, because the substitution is a similarity
+    transform: whatever `1/c` inflates, `c` deflates again on the way out, and a
+    common scale factor does not move a floating-point relative error. Measured at
+    the decay floor, `T=256`, against the sequential form: chunk 16 gives 2.0e-07 and
+    chunk 64 gives 2.7e-07, both fp32 noise, while the recurrence runs 3.9x faster
+    (3.771 ms -> 0.913 ms). 16 was the conservative reading of the same bound.
+    """
+    batch, seq_len, num_heads, head_dim = r.shape
+    dtype = r.dtype
+    r, w_log, k, v, kk, a = (t.to(compute_dtype) for t in (r, w_log, k, v, kk, a))
+    state = state.to(compute_dtype)
+    # Everything that does not touch the carried state is computed for ALL chunks at
+    # once. Only the state recurrence is serial, and it was pulling the rest of the
+    # arithmetic into the Python loop with it: at T=256 and chunk 64 that was four
+    # iterations of ten-odd launches each, plus three constant matrices rebuilt every
+    # iteration. The loop below now does the state-dependent terms and nothing else.
+    # Never pad past the sequence: a 16-token prefill grouped into 64-wide chunks does
+    # four times the arithmetic for the same answer, which cost 1x16 nine points before
+    # this line existed. The old loop avoided it by shortening its last chunk.
+    chunk_size = min(chunk_size, seq_len)
+    # Refuse rather than return NaN. `1/c` reaches `e^(e^-0.5 * chunk_size)` at the
+    # decay floor, so the widest chunk this dtype can carry is derived from the dtype
+    # rather than written down: 146 in fp32, 12 in fp16. A caller who raises
+    # `chunk_size` for a longer prefill gets an error naming the limit instead of
+    # all-NaN logits several layers later, which is what the previous version did.
+    widest = int(math.log(torch.finfo(compute_dtype).max) / _INV_SQRT_E)
+    if chunk_size > widest:
+        raise ValueError(
+            f"chunk_size={chunk_size} overflows {compute_dtype}: the running decay's "
+            f"reciprocal grows like e^(e^-0.5 * chunk_size), so the widest chunk that "
+            f"stays finite is {widest}. Lower chunk_size, or raise compute_dtype."
+        )
+    chunks = (seq_len + chunk_size - 1) // chunk_size
+    padded = chunks * chunk_size
+    if padded != seq_len:
+        pad = (0, 0, 0, 0, 0, padded - seq_len)
+        r, k, v, kk, a = (torch.nn.functional.pad(t, pad) for t in (r, k, v, kk, a))
+        # A pad step must be the identity for the recurrence: decay 1, nothing added.
+        w_log = torch.nn.functional.pad(w_log, pad)
+    grouped = lambda t: t.reshape(batch, chunks, chunk_size, num_heads, head_dim)
+    rg, kg, vg, kkg, ag, wg = (grouped(t) for t in (r, k, v, kk, a, w_log))
+
+    # `cumprod(exp(x))` and `exp(cumsum(x))` are the same function, and the running
+    # decay is only ever needed in the exponentiated form, so the sum is taken first.
+    # Two reasons to prefer it. Numerically, a cumulative product of up-to-1 factors
+    # underflows the further it runs while the equivalent sum of logs does not, and
+    # `c_prev` becomes a subtraction instead of a division by a possibly-tiny `w_c`.
+    # For export, `aten.cumprod` decomposes into a graph carrying a scalar tensor
+    # constant; `aten.lift_fresh` on that constant decomposes to `aten.alias`, which
+    # functionalization rewrites to the in-place `aten.detach_`, and aot_autograd
+    # rejects the result -- which is what made every ONNX export subtest fail.
+    cum = torch.cumsum(wg, dim=2)
+    c = torch.exp(cum)
+    c_prev = torch.exp(cum - wg)
+    bg = kkg * ag
+    b_t, k_t = bg / c, kg / c
+    q_t, r_t = kkg * c_prev, rg * c
+
+    span = chunk_size
+    tri = torch.ones(span, span, device=r.device, dtype=r.dtype).tril(-1)
+    causal = torch.ones(span, span, device=r.device, dtype=r.dtype).tril(0)
+
+    # [batch, chunks, heads, span, span] -- one launch each instead of one per chunk.
+    qb = torch.einsum("bcthn,bcshn->bchts", q_t, b_t) * tri
+    qk = torch.einsum("bcthn,bcshn->bchts", q_t, k_t) * tri
+    rk = torch.einsum("bcthn,bcshn->bchts", r_t, k_t) * causal
+    rb = torch.einsum("bcthn,bcshn->bchts", r_t, b_t) * causal
+    lhs_inv = _unit_lower_triangular_inverse(qb)
+    qkv = torch.einsum("bchts,bcshv->bchtv", qk, vg)
+    rkv = torch.einsum("bchts,bcshv->bchtv", rk, vg)
+    c_last = c[:, :, -1].unsqueeze(-1)
+
+    # `k~^T v` carries no state either, so it joins the batched half above; only the
+    # two products that read the carried state and the solve that depends on them are
+    # genuinely serial.
+    kv = torch.einsum("bcthn,bcthv->bchnv", k_t, vg)
+    # Both state products read the SAME state, so they are one matmul over a stacked
+    # token axis rather than two -- the loop is short and launch-bound, and this is one
+    # launch per chunk instead of two.
+    qr = torch.cat([q_t, r_t], dim=2)
+
+    outputs = []
+    for i in range(chunks):
+        qr_s = torch.einsum("bthn,bhnv->bhtv", qr[:, i], state)
+        rhs = qr_s[:, :, :chunk_size] + qkv[:, i]
+        u = lhs_inv[:, i] @ rhs
+        out_c = (
+            qr_s[:, :, chunk_size:]
+            + rkv[:, i]
+            - torch.einsum("bhts,bhsv->bhtv", rb[:, i], u)
+        )
+        outputs.append(out_c.permute(0, 2, 1, 3))
+        state = c_last[:, i] * (
+            state + kv[:, i] - torch.einsum("bthn,bhtv->bhnv", b_t[:, i], u)
+        )
+
+    # The chunked loop pads the sequence up to a multiple of the chunk size and
+    # computes in a wider dtype, so both the trailing slice and the cast back are
+    # no-ops whenever the input needed neither. A no-op slice or cast still returns
+    # a view, ONNX export turns that view into `aten.alias`, functionalization
+    # rewrites the alias to the in-place `aten.detach_`, and aot_autograd rejects
+    # the graph. Doing each step only when it changes something keeps the common
+    # path free of views without changing any value.
+    out = torch.cat(outputs, dim=1)
+    if out.shape[1] != seq_len:
+        out = out[:, :seq_len]
+    if out.dtype != dtype:
+        out = out.to(dtype)
+    return out, state
+
+
+def rwkv7_eager(
+    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_q: torch.Tensor | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Portable WKV: the sequential step for a single token, chunk-parallel otherwise.
+
+    With `cu_seq_lens_q` 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
+    would resume from. This is the reference behaviour, not the fast path: a
+    fused varlen kernel does the same segments in one launch, which is the shape
+    of optimised implementation the hub-kernels mechanism can substitute in.
+    """
+    if cu_seq_lens_q is not None:
+        bounds = cu_seq_lens_q.tolist()
+        outputs = []
+        for start, stop in zip(bounds[:-1], bounds[1:]):
+            if stop <= start:
+                continue
+            segment_state = torch.zeros_like(state)
+            out, state = rwkv7_eager(
+                r[:, start:stop],
+                w_log[:, start:stop],
+                k[:, start:stop],
+                v[:, start:stop],
+                kk[:, start:stop],
+                a[:, start:stop],
+                segment_state,
+            )
+            outputs.append(out)
+        return torch.cat(outputs, dim=1), state
+
+    kernel = rwkv7_recurrent if r.shape[1] == 1 else rwkv7_chunked
+    return kernel(r, w_log, k, v, kk, a, state)
+
+
+# ---- Triton decode specialisation (defined only when Triton imports) ----
+
+if _HAVE_TRITON:
+
+    @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)
+
+else:
+    fused_wkv_one = None
+
+
+def _use_kernel(r, cu_seq_lens) -> bool:
+    """Whether the fused step applies, as a predicate rather than inline conditions.
+
+    Separated out because it is the part worth testing and the part that cannot be
+    tested through `rwkv7_fused` on a machine without CUDA: there every case falls back
+    for the same reason, so a parametrised test over prefill/packed/width passes even
+    with the other conditions deleted. Checked that by deleting them.
+
+    `head_dim` a power of two is part of the kernel's applicability domain, not a
+    detail: `tl.arange` requires it, and without that term a width like 192 reached
+    Triton and came back as a CompilationError raised from inside the kernel rather
+    than as a fallback.
+    """
+    width = r.shape[-1]
+    return bool(
+        r.shape[1] == 1  # decode only; prefill has its own chunked form
+        and cu_seq_lens is None  # a packed row is many sequences, not one step
+        and r.is_cuda
+        and _HAVE_TRITON
+        and width & (width - 1) == 0
+    )
+
+
+def rwkv7_wkv(
+    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_q: torch.Tensor | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """The WKV recurrence with the decode step fused; reference semantics everywhere.
+
+    Drop-in for the `transformers` RWKV-7 `rwkv7_eager` contract: same tensor
+    layout ([batch, seq_len, num_heads, head_dim] vectors, [batch, num_heads,
+    head_dim, head_dim] state), same `(output, new_state)` return, `cu_seq_lens_q`
+    packed rows honoured by the reference path.
+    """
+    if _use_kernel(r, cu_seq_lens_q):
+        return fused_wkv_one(r, w_log, k, v, kk, a, state), state
+    return rwkv7_eager(r, w_log, k, v, kk, a, state, cu_seq_lens_q=cu_seq_lens_q)
+
+
+__all__ = [
+    "fused_wkv_one",
+    "rwkv7_chunked",
+    "rwkv7_eager",
+    "rwkv7_recurrent",
+    "rwkv7_wkv",
+]