rwkv7-kernel / tests /test_sparse_cmix.py

Add regression tests for the wrapped-value-module gradient bug

cf6a52e verified 3 months ago

17.8 kB

"""The fused projection must agree with the dense one it replaces.
Skipping a zero row is exact, so the only difference allowed is the fp32 atomicAdd
combine's arrival order — ~1 ULP, the split-K GEMM class. The tests therefore compare
against a dense fp32 reference rather than against fp16 `linear`, and check that the
fused path is no further from the truth than the dense path is.
"""

import pytest
import torch

cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")


def _acc(hidden, device="cuda"):
    """The caller-owned fp32 accumulator the op now writes into, one per module."""
    return torch.zeros(hidden, device=device, dtype=torch.float32)


def _sparse_activation(inter, density, device, generator):
    """A squared-ReLU activation: mostly exact zeros, the rest strictly positive."""
    raw = torch.randn(inter, device=device, generator=generator, dtype=torch.float16)
    keep = torch.rand(inter, device=device, generator=generator) < density
    return torch.where(keep, raw.abs(), torch.zeros_like(raw))


@cuda
@pytest.mark.parametrize("density", [0.0, 0.0667, 0.5, 1.0])
def test_matches_dense_reference(density):
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    torch.manual_seed(0)
    g = torch.Generator(device="cuda").manual_seed(0)
    hidden, inter = 512, 1024
    weight = (torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16) * 0.02)
    act = _sparse_activation(inter, density, "cuda", g)

    got = torch.ops.rwkv7_kernel.sparse_cmix(act, _tile_value_weight(weight), _acc(hidden))
    dense = torch.nn.functional.linear(act, weight)
    truth = torch.nn.functional.linear(act.float(), weight.float())

    fused_err = (got.float().squeeze(0) - truth).abs().max()
    dense_err = (dense.float() - truth).abs().max()
    # Not "fused == dense" -- neither is exact in fp16. The bar is that swapping in
    # the kernel does not move the answer further from the fp32 truth than the
    # implementation it replaces already is.
    assert fused_err <= dense_err * 1.5 + 1e-4, (fused_err.item(), dense_err.item())


@cuda
def test_all_zero_activation_gives_zero():
    """The density-0 case has no rows to accumulate; the accumulator must still be
    clean, which is only true if the previous call re-zeroed it."""
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    g = torch.Generator(device="cuda").manual_seed(1)
    hidden, inter = 512, 1024
    weight = torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16)
    tiled = _tile_value_weight(weight)
    acc = _acc(hidden)

    torch.ops.rwkv7_kernel.sparse_cmix(_sparse_activation(inter, 0.5, "cuda", g), tiled, acc)
    zeros = torch.zeros(inter, device="cuda", dtype=torch.float16)
    got = torch.ops.rwkv7_kernel.sparse_cmix(zeros, tiled, acc)
    assert torch.equal(got, torch.zeros_like(got)), got.abs().max().item()


@cuda
def test_repeated_calls_do_not_accumulate():
    """The persistent accumulator is reused across layers and steps. If the finalize
    pass failed to re-zero it, the second identical call would return roughly twice
    the first -- which no correctness test on a single call would notice."""
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    g = torch.Generator(device="cuda").manual_seed(2)
    hidden, inter = 512, 1024
    weight = (torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16) * 0.02)
    tiled = _tile_value_weight(weight)
    acc = _acc(hidden)
    act = _sparse_activation(inter, 0.0667, "cuda", g)

    first = torch.ops.rwkv7_kernel.sparse_cmix(act, tiled, acc).clone()
    for _ in range(5):
        again = torch.ops.rwkv7_kernel.sparse_cmix(act, tiled, acc)
    torch.testing.assert_close(first, again, rtol=0, atol=0)


@cuda
def test_tiled_weight_follows_weight_updates():
    """A cached tile that survives a weight edit projects with the old weights and
    says nothing about it."""
    from rwkv7_kernel.sparse_ffn import _tiled_weight

    hidden, inter = 512, 1024
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()

    before = _tiled_weight(module).clone()
    with torch.no_grad():
        module.value.weight.mul_(2.0)
    after = _tiled_weight(module)

    torch.testing.assert_close(after, before * 2, rtol=1e-3, atol=1e-3)


@cuda
def test_two_modules_do_not_share_an_accumulator():
    """The bug that made the accumulator caller-owned.
    It used to be a process-wide map keyed on (device, hidden), so two modules of the
    same width shared one buffer. On a single stream that is still correct -- kernel
    completion orders the reuse -- and off it the partial sums interleave and the
    logits come back wrong with no NaN and no error. Two modules of equal width must
    now hold two distinct buffers; the same object would put the old hazard back.
    """
    from rwkv7_kernel.sparse_ffn import _tiled_weight

    hidden, inter = 512, 1024
    modules = []
    for _ in range(2):
        m = torch.nn.Module()
        m.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
        _tiled_weight(m)
        modules.append(m)

    a, b = (m._rwkv_hf_accumulator for m in modules)
    assert a.data_ptr() != b.data_ptr()
    assert a.shape == (hidden,) and a.dtype == torch.float32


@cuda
def test_mismatched_devices_are_refused():
    """A device_map split would otherwise launch on one GPU with another's pointers."""
    if torch.cuda.device_count() < 2:
        pytest.skip("needs two CUDA devices")
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    hidden, inter = 512, 1024
    weight = torch.randn(hidden, inter, device="cuda:0", dtype=torch.float16) * 0.02
    act = torch.zeros(inter, device="cuda:1", dtype=torch.float16)
    with pytest.raises(RuntimeError, match="share a device"):
        torch.ops.rwkv7_kernel.sparse_cmix(act, _tile_value_weight(weight), _acc(hidden, "cuda:0"))


@cuda
def test_grad_enabled_forward_is_refused_rather_than_silently_ungraded():
    """There is no backward, so a grad-enabled forward used to return a value whose
    gradient never arrives -- `.backward()` succeeds and the input grad is None."""
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 512, 1024
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
    act = torch.zeros(1, 1, inter, device="cuda", dtype=torch.float16, requires_grad=True)

    with pytest.raises(RuntimeError, match="no backward"):
        sparse_project(module, act)

    with torch.no_grad():  # the documented way through
        sparse_project(module, act)


def test_register_with_a_model_it_cannot_pre_tile_raises():
    """Passing a model means "build the tiles now". Finding none used to return
    quietly, costing the caller the CUDA-graph cliff the eager build exists to avoid
    with nothing said. Runs on CPU: the argument check now fires before load_extension()."""
    from rwkv7_kernel import sparse_ffn

    class NotAnRwkv(torch.nn.Module):
        pass

    with pytest.raises(ValueError, match="no blocks to pre-tile"):
        sparse_ffn.register(NotAnRwkv())


def test_invalidate_forces_a_rebuild():
    """The escape hatch for the edits the fingerprint cannot see."""
    from rwkv7_kernel.sparse_ffn import invalidate

    module = torch.nn.Module()
    module._rwkv_hf_tiled_fp = ("pretend a tile was built",)
    invalidate(module)
    assert module._rwkv_hf_tiled_fp is None


def test_register_that_raises_leaves_the_class_untouched():
    """The raise above exists to guarantee this, and nothing was checking it.
    `register` promises that a rejected call "leaves nothing half-applied". The
    existing test asserts the exception and stops there, so the promise the exception
    is there to keep was never exercised: an implementation that patched the class
    first and validated second would pass it.
    """
    from rwkv7_kernel import sparse_ffn

    # Guard on the module the import below needs; `transformers` alone is satisfied by a
    # stock install that has no RWKV-7, and the test then errors instead of skipping.
    pytest.importorskip("transformers.models.rwkv7")
    from transformers.models.rwkv7.modeling_rwkv7 import Rwkv7FeedForward

    before = Rwkv7FeedForward._project
    with pytest.raises(ValueError, match="no blocks to pre-tile"):
        sparse_ffn.register(torch.nn.Module())

    assert Rwkv7FeedForward._project is before
    assert sparse_ffn._ORIGINAL_PROJECT is None


def test_unregister_restores_the_projection_and_is_idempotent():
    """`unregister` had no test at all, on either of its two paths.
    Driven by setting `_ORIGINAL_PROJECT` directly rather than by calling `register`,
    which would pull in the CUDA extension; the restore path is plain Python and this
    is how `test_invalidate_forces_a_rebuild` already exercises its own state.
    """
    from rwkv7_kernel import sparse_ffn

    pytest.importorskip("transformers.models.rwkv7")
    from transformers.models.rwkv7.modeling_rwkv7 import Rwkv7FeedForward

    original = Rwkv7FeedForward._project
    try:
        sparse_ffn._ORIGINAL_PROJECT = original
        Rwkv7FeedForward._project = sparse_ffn.sparse_project

        sparse_ffn.unregister()
        assert Rwkv7FeedForward._project is original
        assert sparse_ffn._ORIGINAL_PROJECT is None

        # Calling it again must not re-restore a stale function or raise. This is the
        # branch that returns early, and it is the one a caller hits by unregistering
        # in a `finally` that also runs on the path where register never happened.
        sparse_ffn.unregister()
        assert Rwkv7FeedForward._project is original
    finally:
        Rwkv7FeedForward._project = original
        sparse_ffn._ORIGINAL_PROJECT = None


def test_a_non_conforming_input_takes_the_dense_path_and_builds_no_tile():
    """The fallback branch was uncovered, and it is the one most callers hit first.
    Bit-exact against `linear` rather than close to it: the branch does not
    approximate the dense projection, it *is* the dense projection, so any tolerance
    would also pass an implementation that quietly ran the kernel. The tile assertion
    is what actually pins the branch -- a conforming call would have built one.
    """
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 256, 512
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False)
    act = torch.randn(1, 1, inter)  # CPU, fp32: non-conforming on two counts

    got = sparse_project(module, act)

    torch.testing.assert_close(got, torch.nn.functional.linear(act, module.value.weight), rtol=0, atol=0)
    assert getattr(module, "_rwkv_hf_tiled_fp", None) is None
    assert not hasattr(module, "_rwkv_hf_tiled")


class _FakeAdapterWrapper(torch.nn.Module):
    """Minimal stand-in for what PEFT's `lora.Linear` actually looks like from the
    outside: a `.weight` attribute that proxies to a frozen base layer (so shape
    checks upstream still work), plus a `forward()` that adds something the base
    layer alone does not compute. Not importing peft here on purpose -- the bug
    this guards is about *any* module with its own forward() sitting where a bare
    nn.Linear was assumed, not specifically about LoRA's internals.
    """

    def __init__(self, base: torch.nn.Linear, delta: torch.nn.Linear):
        super().__init__()
        self.base = base
        self.delta = delta

    @property
    def weight(self):
        return self.base.weight

    def forward(self, x):
        return self.base(x) + self.delta(x)


def test_a_wrapped_value_module_goes_through_its_own_forward_not_raw_weight():
    """Regression: an earlier version of both fallback branches did
    `F.linear(activated, module.value.weight)`, which reaches past whatever
    `module.value` actually is straight to the frozen base weight. That is exactly
    what `Rwkv7FeedForward._project` (`return self.value(inner)`) does not do, so it
    silently stopped being a drop-in the moment `module.value` was anything other
    than a bare `nn.Linear` -- confirmed against a real PEFT LoRA adapter on
    `ffn.value`: `lora_A`/`lora_B` came back with `grad=None` (never entered the
    graph at all) on every training-shaped call, because `_conforms` requires
    single-token input and a normal training forward never qualifies. This uses a
    hand-rolled wrapper instead of peft so the test has no extra dependency, but the
    property under test -- "does the module's own forward() get called" -- is the
    same one that broke.
    """
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 256, 512
    base = torch.nn.Linear(inter, hidden, bias=False)
    base.weight.requires_grad_(False)  # PEFT freezes the base layer; mirror that
    delta = torch.nn.Linear(inter, hidden, bias=False)
    module = torch.nn.Module()
    module.value = _FakeAdapterWrapper(base, delta)

    act = torch.randn(1, 1, inter, requires_grad=True)  # non-conforming: CPU + fp32

    got = sparse_project(module, act)
    expected = module.value(act)  # what the original _project would have produced

    torch.testing.assert_close(got, expected, rtol=0, atol=0)
    assert not torch.allclose(got, torch.nn.functional.linear(act, base.weight)), (
        "the fallback produced the same answer as skipping the wrapper entirely -- "
        "delta's contribution did not survive, so this test is not actually checking "
        "what it claims to"
    )

    got.sum().backward()
    assert delta.weight.grad is not None and delta.weight.grad.abs().sum() > 0, (
        "the wrapper's own trainable parameter got no gradient -- the bug this "
        "guards against is back"
    )


@cuda
def test_conforming_shape_also_declines_a_wrapped_value_module():
    """The fused-kernel path itself reads `module.value.weight` directly too (for
    tiling), which for a wrapped module is the unadapted base weight -- silently
    wrong under `torch.no_grad()` single-token inference with an unmerged adapter,
    not just a training-time gradient bug. A wrapped module must decline the fused
    path even when the shape/dtype would otherwise conform, and take the module's
    own forward() instead, same as the non-conforming branch."""
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 256, 512
    base = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
    base.weight.requires_grad_(False)
    delta = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
    module = torch.nn.Module()
    module.value = _FakeAdapterWrapper(base, delta)

    act = torch.zeros(1, 1, inter, device="cuda", dtype=torch.float16)  # conforming shape/dtype

    with torch.no_grad():
        got = sparse_project(module, act)
        expected = module.value(act)

    torch.testing.assert_close(got, expected, rtol=0, atol=0)
    assert getattr(module, "_rwkv_hf_tiled_fp", None) is None, (
        "a tile got built for a wrapped module -- the fused path should have "
        "declined before ever reaching _tiled_weight"
    )


@cuda
@pytest.mark.parametrize(
    "reason,shape,dtype",
    [
        ("more than one token", (1, 2), torch.float16),
        ("not fp16", (1, 1), torch.float32),
    ],
)
def test_each_non_conforming_reason_falls_back_on_its_own(reason, shape, dtype):
    """On CPU every input is non-conforming for the same reason, so the individual
    conditions can only be separated on a device where the others are satisfied."""
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 512, 1024
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False).to("cuda", dtype)
    act = torch.randn(*shape, inter, device="cuda", dtype=dtype)

    with torch.no_grad():
        got = sparse_project(module, act)
        want = torch.nn.functional.linear(act, module.value.weight)

    torch.testing.assert_close(got, want, rtol=0, atol=0), reason
    assert getattr(module, "_rwkv_hf_tiled_fp", None) is None, reason


def test_a_non_default_tile_size_is_read_and_the_layout_still_inverts(monkeypatch):
    """No tile size other than the default was covered anywhere.
    This pins the Python half only: `_tiles` reading the environment at call time,
    and `_tile_value_weight` staying a pure permutation at that size. It deliberately
    does *not* claim the kernel agrees -- the tile shape is a compile-time constant
    on that side, so an env var changed after the extension is built gives the two
    halves different constants, which is exactly why `_tiles` reads late. Checking
    kernel agreement at a second size means rebuilding the extension, which belongs
    in a build matrix rather than here.
    """
    from rwkv7_kernel.sparse_ffn import _tile_value_weight, _tiles

    monkeypatch.setenv("RWKV_HF_KERNELS_C_TILE", "128")
    tile_h, tile_i = _tiles()
    assert (tile_h, tile_i) == (128, 64)

    hidden, inter = 256, 128
    weight = torch.randn(hidden, inter)
    tiled = _tile_value_weight(weight)
    assert tiled.shape == (inter // tile_i, hidden // tile_h, tile_i, tile_h)

    # A permutation loses nothing: undo it and the original must come back exactly.
    back = tiled.permute(0, 2, 1, 3).reshape(inter, hidden).t()
    torch.testing.assert_close(back, weight, rtol=0, atol=0)

	"""The fused projection must agree with the dense one it replaces.

	Skipping a zero row is exact, so the only difference allowed is the fp32 atomicAdd
	combine's arrival order — ~1 ULP, the split-K GEMM class. The tests therefore compare
	against a dense fp32 reference rather than against fp16 `linear`, and check that the
	fused path is no further from the truth than the dense path is.
	"""

	import pytest
	import torch

	cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")


	def _acc(hidden, device="cuda"):
	"""The caller-owned fp32 accumulator the op now writes into, one per module."""
	return torch.zeros(hidden, device=device, dtype=torch.float32)


	def _sparse_activation(inter, density, device, generator):
	"""A squared-ReLU activation: mostly exact zeros, the rest strictly positive."""
	raw = torch.randn(inter, device=device, generator=generator, dtype=torch.float16)
	keep = torch.rand(inter, device=device, generator=generator) < density
	return torch.where(keep, raw.abs(), torch.zeros_like(raw))


	@cuda
	@pytest.mark.parametrize("density", [0.0, 0.0667, 0.5, 1.0])
	def test_matches_dense_reference(density):
	from rwkv7_kernel._ext import load_extension
	from rwkv7_kernel.sparse_ffn import _tile_value_weight

	load_extension()
	torch.manual_seed(0)
	g = torch.Generator(device="cuda").manual_seed(0)
	hidden, inter = 512, 1024
	weight = (torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16) * 0.02)
	act = _sparse_activation(inter, density, "cuda", g)

	got = torch.ops.rwkv7_kernel.sparse_cmix(act, _tile_value_weight(weight), _acc(hidden))
	dense = torch.nn.functional.linear(act, weight)
	truth = torch.nn.functional.linear(act.float(), weight.float())

	fused_err = (got.float().squeeze(0) - truth).abs().max()
	dense_err = (dense.float() - truth).abs().max()
	# Not "fused == dense" -- neither is exact in fp16. The bar is that swapping in
	# the kernel does not move the answer further from the fp32 truth than the
	# implementation it replaces already is.
	assert fused_err <= dense_err * 1.5 + 1e-4, (fused_err.item(), dense_err.item())


	@cuda
	def test_all_zero_activation_gives_zero():
	"""The density-0 case has no rows to accumulate; the accumulator must still be
	clean, which is only true if the previous call re-zeroed it."""
	from rwkv7_kernel._ext import load_extension
	from rwkv7_kernel.sparse_ffn import _tile_value_weight

	load_extension()
	g = torch.Generator(device="cuda").manual_seed(1)
	hidden, inter = 512, 1024
	weight = torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16)
	tiled = _tile_value_weight(weight)
	acc = _acc(hidden)

	torch.ops.rwkv7_kernel.sparse_cmix(_sparse_activation(inter, 0.5, "cuda", g), tiled, acc)
	zeros = torch.zeros(inter, device="cuda", dtype=torch.float16)
	got = torch.ops.rwkv7_kernel.sparse_cmix(zeros, tiled, acc)
	assert torch.equal(got, torch.zeros_like(got)), got.abs().max().item()


	@cuda
	def test_repeated_calls_do_not_accumulate():
	"""The persistent accumulator is reused across layers and steps. If the finalize
	pass failed to re-zero it, the second identical call would return roughly twice
	the first -- which no correctness test on a single call would notice."""
	from rwkv7_kernel._ext import load_extension
	from rwkv7_kernel.sparse_ffn import _tile_value_weight

	load_extension()
	g = torch.Generator(device="cuda").manual_seed(2)
	hidden, inter = 512, 1024
	weight = (torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16) * 0.02)
	tiled = _tile_value_weight(weight)
	acc = _acc(hidden)
	act = _sparse_activation(inter, 0.0667, "cuda", g)

	first = torch.ops.rwkv7_kernel.sparse_cmix(act, tiled, acc).clone()
	for _ in range(5):
	again = torch.ops.rwkv7_kernel.sparse_cmix(act, tiled, acc)
	torch.testing.assert_close(first, again, rtol=0, atol=0)


	@cuda
	def test_tiled_weight_follows_weight_updates():
	"""A cached tile that survives a weight edit projects with the old weights and
	says nothing about it."""
	from rwkv7_kernel.sparse_ffn import _tiled_weight

	hidden, inter = 512, 1024
	module = torch.nn.Module()
	module.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()

	before = _tiled_weight(module).clone()
	with torch.no_grad():
	module.value.weight.mul_(2.0)
	after = _tiled_weight(module)

	torch.testing.assert_close(after, before * 2, rtol=1e-3, atol=1e-3)


	@cuda
	def test_two_modules_do_not_share_an_accumulator():
	"""The bug that made the accumulator caller-owned.

	It used to be a process-wide map keyed on (device, hidden), so two modules of the
	same width shared one buffer. On a single stream that is still correct -- kernel
	completion orders the reuse -- and off it the partial sums interleave and the
	logits come back wrong with no NaN and no error. Two modules of equal width must
	now hold two distinct buffers; the same object would put the old hazard back.
	"""
	from rwkv7_kernel.sparse_ffn import _tiled_weight

	hidden, inter = 512, 1024
	modules = []
	for _ in range(2):
	m = torch.nn.Module()
	m.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
	_tiled_weight(m)
	modules.append(m)

	a, b = (m._rwkv_hf_accumulator for m in modules)
	assert a.data_ptr() != b.data_ptr()
	assert a.shape == (hidden,) and a.dtype == torch.float32


	@cuda
	def test_mismatched_devices_are_refused():
	"""A device_map split would otherwise launch on one GPU with another's pointers."""
	if torch.cuda.device_count() < 2:
	pytest.skip("needs two CUDA devices")
	from rwkv7_kernel._ext import load_extension
	from rwkv7_kernel.sparse_ffn import _tile_value_weight

	load_extension()
	hidden, inter = 512, 1024
	weight = torch.randn(hidden, inter, device="cuda:0", dtype=torch.float16) * 0.02
	act = torch.zeros(inter, device="cuda:1", dtype=torch.float16)
	with pytest.raises(RuntimeError, match="share a device"):
	torch.ops.rwkv7_kernel.sparse_cmix(act, _tile_value_weight(weight), _acc(hidden, "cuda:0"))


	@cuda
	def test_grad_enabled_forward_is_refused_rather_than_silently_ungraded():
	"""There is no backward, so a grad-enabled forward used to return a value whose
	gradient never arrives -- `.backward()` succeeds and the input grad is None."""
	from rwkv7_kernel.sparse_ffn import sparse_project

	hidden, inter = 512, 1024
	module = torch.nn.Module()
	module.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
	act = torch.zeros(1, 1, inter, device="cuda", dtype=torch.float16, requires_grad=True)

	with pytest.raises(RuntimeError, match="no backward"):
	sparse_project(module, act)

	with torch.no_grad(): # the documented way through
	sparse_project(module, act)


	def test_register_with_a_model_it_cannot_pre_tile_raises():
	"""Passing a model means "build the tiles now". Finding none used to return
	quietly, costing the caller the CUDA-graph cliff the eager build exists to avoid
	with nothing said. Runs on CPU: the argument check now fires before load_extension()."""
	from rwkv7_kernel import sparse_ffn

	class NotAnRwkv(torch.nn.Module):
	pass

	with pytest.raises(ValueError, match="no blocks to pre-tile"):
	sparse_ffn.register(NotAnRwkv())


	def test_invalidate_forces_a_rebuild():
	"""The escape hatch for the edits the fingerprint cannot see."""
	from rwkv7_kernel.sparse_ffn import invalidate

	module = torch.nn.Module()
	module._rwkv_hf_tiled_fp = ("pretend a tile was built",)
	invalidate(module)
	assert module._rwkv_hf_tiled_fp is None


	def test_register_that_raises_leaves_the_class_untouched():
	"""The raise above exists to guarantee this, and nothing was checking it.

	`register` promises that a rejected call "leaves nothing half-applied". The
	existing test asserts the exception and stops there, so the promise the exception
	is there to keep was never exercised: an implementation that patched the class
	first and validated second would pass it.
	"""
	from rwkv7_kernel import sparse_ffn

	# Guard on the module the import below needs; `transformers` alone is satisfied by a
	# stock install that has no RWKV-7, and the test then errors instead of skipping.
	pytest.importorskip("transformers.models.rwkv7")
	from transformers.models.rwkv7.modeling_rwkv7 import Rwkv7FeedForward

	before = Rwkv7FeedForward._project
	with pytest.raises(ValueError, match="no blocks to pre-tile"):
	sparse_ffn.register(torch.nn.Module())

	assert Rwkv7FeedForward._project is before
	assert sparse_ffn._ORIGINAL_PROJECT is None


	def test_unregister_restores_the_projection_and_is_idempotent():
	"""`unregister` had no test at all, on either of its two paths.

	Driven by setting `_ORIGINAL_PROJECT` directly rather than by calling `register`,
	which would pull in the CUDA extension; the restore path is plain Python and this
	is how `test_invalidate_forces_a_rebuild` already exercises its own state.
	"""
	from rwkv7_kernel import sparse_ffn

	pytest.importorskip("transformers.models.rwkv7")
	from transformers.models.rwkv7.modeling_rwkv7 import Rwkv7FeedForward

	original = Rwkv7FeedForward._project
	try:
	sparse_ffn._ORIGINAL_PROJECT = original
	Rwkv7FeedForward._project = sparse_ffn.sparse_project

	sparse_ffn.unregister()
	assert Rwkv7FeedForward._project is original
	assert sparse_ffn._ORIGINAL_PROJECT is None

	# Calling it again must not re-restore a stale function or raise. This is the
	# branch that returns early, and it is the one a caller hits by unregistering
	# in a `finally` that also runs on the path where register never happened.
	sparse_ffn.unregister()
	assert Rwkv7FeedForward._project is original
	finally:
	Rwkv7FeedForward._project = original
	sparse_ffn._ORIGINAL_PROJECT = None


	def test_a_non_conforming_input_takes_the_dense_path_and_builds_no_tile():
	"""The fallback branch was uncovered, and it is the one most callers hit first.

	Bit-exact against `linear` rather than close to it: the branch does not
	approximate the dense projection, it is the dense projection, so any tolerance
	would also pass an implementation that quietly ran the kernel. The tile assertion
	is what actually pins the branch -- a conforming call would have built one.
	"""
	from rwkv7_kernel.sparse_ffn import sparse_project

	hidden, inter = 256, 512
	module = torch.nn.Module()
	module.value = torch.nn.Linear(inter, hidden, bias=False)
	act = torch.randn(1, 1, inter) # CPU, fp32: non-conforming on two counts

	got = sparse_project(module, act)

	torch.testing.assert_close(got, torch.nn.functional.linear(act, module.value.weight), rtol=0, atol=0)
	assert getattr(module, "_rwkv_hf_tiled_fp", None) is None
	assert not hasattr(module, "_rwkv_hf_tiled")


	class _FakeAdapterWrapper(torch.nn.Module):
	"""Minimal stand-in for what PEFT's `lora.Linear` actually looks like from the
	outside: a `.weight` attribute that proxies to a frozen base layer (so shape
	checks upstream still work), plus a `forward()` that adds something the base
	layer alone does not compute. Not importing peft here on purpose -- the bug
	this guards is about any module with its own forward() sitting where a bare
	nn.Linear was assumed, not specifically about LoRA's internals.
	"""

	def __init__(self, base: torch.nn.Linear, delta: torch.nn.Linear):
	super().__init__()
	self.base = base
	self.delta = delta

	@property
	def weight(self):
	return self.base.weight

	def forward(self, x):
	return self.base(x) + self.delta(x)


	def test_a_wrapped_value_module_goes_through_its_own_forward_not_raw_weight():
	"""Regression: an earlier version of both fallback branches did
	`F.linear(activated, module.value.weight)`, which reaches past whatever
	`module.value` actually is straight to the frozen base weight. That is exactly
	what `Rwkv7FeedForward._project` (`return self.value(inner)`) does not do, so it
	silently stopped being a drop-in the moment `module.value` was anything other
	than a bare `nn.Linear` -- confirmed against a real PEFT LoRA adapter on
	`ffn.value`: `lora_A`/`lora_B` came back with `grad=None` (never entered the
	graph at all) on every training-shaped call, because `_conforms` requires
	single-token input and a normal training forward never qualifies. This uses a
	hand-rolled wrapper instead of peft so the test has no extra dependency, but the
	property under test -- "does the module's own forward() get called" -- is the
	same one that broke.
	"""
	from rwkv7_kernel.sparse_ffn import sparse_project

	hidden, inter = 256, 512
	base = torch.nn.Linear(inter, hidden, bias=False)
	base.weight.requires_grad_(False) # PEFT freezes the base layer; mirror that
	delta = torch.nn.Linear(inter, hidden, bias=False)
	module = torch.nn.Module()
	module.value = _FakeAdapterWrapper(base, delta)

	act = torch.randn(1, 1, inter, requires_grad=True) # non-conforming: CPU + fp32

	got = sparse_project(module, act)
	expected = module.value(act) # what the original _project would have produced

	torch.testing.assert_close(got, expected, rtol=0, atol=0)
	assert not torch.allclose(got, torch.nn.functional.linear(act, base.weight)), (
	"the fallback produced the same answer as skipping the wrapper entirely -- "
	"delta's contribution did not survive, so this test is not actually checking "
	"what it claims to"
	)

	got.sum().backward()
	assert delta.weight.grad is not None and delta.weight.grad.abs().sum() > 0, (
	"the wrapper's own trainable parameter got no gradient -- the bug this "
	"guards against is back"
	)


	@cuda
	def test_conforming_shape_also_declines_a_wrapped_value_module():
	"""The fused-kernel path itself reads `module.value.weight` directly too (for
	tiling), which for a wrapped module is the unadapted base weight -- silently
	wrong under `torch.no_grad()` single-token inference with an unmerged adapter,
	not just a training-time gradient bug. A wrapped module must decline the fused
	path even when the shape/dtype would otherwise conform, and take the module's
	own forward() instead, same as the non-conforming branch."""
	from rwkv7_kernel.sparse_ffn import sparse_project

	hidden, inter = 256, 512
	base = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
	base.weight.requires_grad_(False)
	delta = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
	module = torch.nn.Module()
	module.value = _FakeAdapterWrapper(base, delta)

	act = torch.zeros(1, 1, inter, device="cuda", dtype=torch.float16) # conforming shape/dtype

	with torch.no_grad():
	got = sparse_project(module, act)
	expected = module.value(act)

	torch.testing.assert_close(got, expected, rtol=0, atol=0)
	assert getattr(module, "_rwkv_hf_tiled_fp", None) is None, (
	"a tile got built for a wrapped module -- the fused path should have "
	"declined before ever reaching _tiled_weight"
	)


	@cuda
	@pytest.mark.parametrize(
	"reason,shape,dtype",
	[
	("more than one token", (1, 2), torch.float16),
	("not fp16", (1, 1), torch.float32),
	],
	)
	def test_each_non_conforming_reason_falls_back_on_its_own(reason, shape, dtype):
	"""On CPU every input is non-conforming for the same reason, so the individual
	conditions can only be separated on a device where the others are satisfied."""
	from rwkv7_kernel.sparse_ffn import sparse_project

	hidden, inter = 512, 1024
	module = torch.nn.Module()
	module.value = torch.nn.Linear(inter, hidden, bias=False).to("cuda", dtype)
	act = torch.randn(*shape, inter, device="cuda", dtype=dtype)

	with torch.no_grad():
	got = sparse_project(module, act)
	want = torch.nn.functional.linear(act, module.value.weight)

	torch.testing.assert_close(got, want, rtol=0, atol=0), reason
	assert getattr(module, "_rwkv_hf_tiled_fp", None) is None, reason


	def test_a_non_default_tile_size_is_read_and_the_layout_still_inverts(monkeypatch):
	"""No tile size other than the default was covered anywhere.

	This pins the Python half only: `_tiles` reading the environment at call time,
	and `_tile_value_weight` staying a pure permutation at that size. It deliberately
	does not claim the kernel agrees -- the tile shape is a compile-time constant
	on that side, so an env var changed after the extension is built gives the two
	halves different constants, which is exactly why `_tiles` reads late. Checking
	kernel agreement at a second size means rebuilding the extension, which belongs
	in a build matrix rather than here.
	"""
	from rwkv7_kernel.sparse_ffn import _tile_value_weight, _tiles

	monkeypatch.setenv("RWKV_HF_KERNELS_C_TILE", "128")
	tile_h, tile_i = _tiles()
	assert (tile_h, tile_i) == (128, 64)

	hidden, inter = 256, 128
	weight = torch.randn(hidden, inter)
	tiled = _tile_value_weight(weight)
	assert tiled.shape == (inter // tile_i, hidden // tile_h, tile_i, tile_h)

	# A permutation loses nothing: undo it and the original must come back exactly.
	back = tiled.permute(0, 2, 1, 3).reshape(inter, hidden).t()
	torch.testing.assert_close(back, weight, rtol=0, atol=0)