| """Shared test setup. |
| |
| HF's Qwen3.5 / LFM2 models dispatch their causal depthwise conv to the ``causal_conv1d`` |
| hub kernel whenever that package is installed; it is CUDA-only, so the HF reference |
| forwards the CPU tests compare against would fail on CPU tensors. CPU calls take HF's |
| own pure-torch fallback (the same arithmetic); CUDA calls keep the kernel. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import pytest |
|
|
|
|
| def _cpu_fallbacks(monkeypatch) -> None: |
| import torch |
| import torch.nn.functional as F |
|
|
| def conv_fn(x, weight, bias=None, activation=None, **_): |
| out = F.conv1d(x.to(weight.dtype), weight.unsqueeze(1), bias, |
| padding=weight.shape[-1] - 1, groups=x.shape[1])[:, :, :x.shape[-1]] |
| return (F.silu(out) if activation else out).to(x.dtype) |
|
|
| def conv_update(x, conv_state, weight, bias=None, activation=None): |
| state_len = conv_state.shape[-1] |
| cat = torch.cat([conv_state, x], -1).to(weight.dtype) |
| conv_state.copy_(cat[:, :, -state_len:]) |
| out = F.conv1d(cat, weight.unsqueeze(1), bias, groups=x.shape[1])[:, :, -x.shape[-1]:] |
| return (F.silu(out) if activation else out).to(x.dtype) |
|
|
| for mod_name in ("transformers.models.qwen3_5.modeling_qwen3_5", |
| "transformers.models.lfm2.modeling_lfm2"): |
| try: |
| mod = __import__(mod_name, fromlist=["_"]) |
| except ImportError: |
| continue |
| for name, ref in (("causal_conv1d_fn", conv_fn), ("causal_conv1d_update", conv_update)): |
| orig = getattr(mod, name, None) |
| if orig is None: |
| continue |
|
|
| def wrapped(x, *a, _orig=orig, _ref=ref, **kw): |
| return _orig(x, *a, **kw) if x.is_cuda else _ref(x, *a, **kw) |
|
|
| monkeypatch.setattr(mod, name, wrapped) |
|
|
|
|
| @pytest.fixture(autouse=True) |
| def _hf_conv_cpu_fallback(monkeypatch): |
| try: |
| import torch |
| except ImportError: |
| yield |
| return |
| _cpu_fallbacks(monkeypatch) |
| yield |
|
|
