Out-of-range token IDs trigger a CUDA device-side assert instead of a validation error

Description

Passing a token ID outside the model vocabulary to an RWKV-7 model running on CUDA causes a low-level device-side assert. The CUDA context then remains in an invalid state, so the Python/Jupyter process must be restarted.

The runtime currently indexes the embedding matrix directly without validating token IDs first:

x = z["emb.weight"][idx]

This can happen accidentally when a tokenizer incompatible with the checkpoint is used. In my case, it occurred while integrating the 65,536-token RWKV World tokenizer with a Pile checkpoint whose vocabulary size is 50,304.

The tokenizer mismatch was an integration mistake, but the resulting failure mode could be handled safely by the runtime.

Environment

• rwkv==0.8.31
• Python 3.12
• PyTorch 2.8.0+cu128
• CUDA 12.8
• NVIDIA L40S
• Ubuntu 24.04
• Model: RWKV-x070-Pile-421M-20241127-ctx4096
• Strategy: cuda fp16

Environment variables:

os.environ["RWKV_V7_ON"] = "1"
os.environ["RWKV_JIT_ON"] = "1"
os.environ["RWKV_CUDA_ON"] = "1"

Minimal reproduction

import os

os.environ["RWKV_V7_ON"] = "1"
os.environ["RWKV_JIT_ON"] = "1"
os.environ["RWKV_CUDA_ON"] = "1"

from rwkv.model import RWKV

model = RWKV(
    model="/path/to/RWKV-x070-Pile-421M-20241127-ctx4096",
    strategy="cuda fp16",
)

print("Vocabulary size:", model.args.vocab_size)

# The first invalid token ID:
invalid_token = int(model.args.vocab_size)

model.forward([invalid_token], None)

Actual behavior

The invalid index reaches the CUDA embedding lookup and triggers an assertion similar to:

/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16:
vectorized_gather_kernel:
Assertion `ind >= 0 && ind < ind_dim_size &&
"vectorized gather kernel index out of bounds"` failed.

RuntimeError: CUDA error: device-side assert triggered

Because CUDA operations are asynchronous, the Python traceback may point to an unrelated later operation inside the RWKV forward pass.

After the assertion, subsequent CUDA operations also fail, and the Jupyter kernel or Python process must be restarted.

Expected behavior

The runtime should validate token IDs before launching any CUDA operation and raise a descriptive Python exception, for example:

ValueError: Token ID 50304 is outside the model vocabulary range [0, 50303].
This may indicate a tokenizer/checkpoint mismatch.

The CUDA context should remain usable after the exception.

Suggested fix

Validate scalar and list inputs at the beginning of RWKV.forward, before calling forward_one or forward_seq:

def forward(self, idx, state, full_output=False):
    token_ids = idx if isinstance(idx, list) else [idx]

    if not token_ids:
        raise ValueError("Token sequence must not be empty.")

    minimum = min(token_ids)
    maximum = max(token_ids)
    vocab_size = int(self.args.vocab_size)

    if minimum < 0 or maximum >= vocab_size:
        raise ValueError(
            f"Token IDs are outside the valid range: "
            f"received [{minimum}, {maximum}], "
            f"expected [0, {vocab_size - 1}]. "
            "Check whether the tokenizer matches the checkpoint."
        )

    # Existing forward implementation...

This check is inexpensive compared with model inference and would prevent an otherwise unrecoverable CUDA failure.

Additional context

The official documentation correctly distinguishes between:

• 20B_tokenizer.json for Pile checkpoints;
• rwkv_vocab_v20230424 for World/G-series checkpoints.

This issue is therefore not about automatically selecting a tokenizer. It is about failing safely and providing a useful diagnostic when incompatible token IDs are supplied.