Add XLSR-MamBo arena wrapper (PyTorch + TensorRT drop-in, vendored mamba_ssm subset)

6beaffa verified 2 days ago

13.9 kB

"""XLSR-MamBo network (MamBo-3: Hydra-N3, 5 repeats), ported from
https://github.com/saki-ciallo/XLSR-MamBo (MIT) @ f429bb1:
- model_build.F4_XLSR300M_General            -> XLSRMamBo
- modules/step1_feature_ptm.WLSR300M_Extractor + RMS_InputProj
- modules/step2_model_blocks.{Hydra_Block, MLP_block, MHA_Block, MLP_SwiGLU,
  Mix_n_Hydra_1MLP_1MHA_1MLP}
- modules/step3_pooling_outputs.GatedAttentionPool
- modules/hydra (Hydra mixer, vendored in ./_hydra)
- train_config: task xlsr_Hydra = Mamba2Config(d_model=128, d_state=64, d_conv=7,
  expand=2, mha_num_heads=4, mha_mlp_dim=0, mha_d_conv=4, residual_in_fp32=True,
  fused_add_norm=True, use_mem_eff_path=True); block_layers=5, n_mamba=3.
Parameter/module names are kept identical to the source so the released
state_dict loads strictly.
Two numerically-equivalent backends:
  * ``pure_torch=False`` (default): the source path -- Triton kernels from the
    vendored mamba_ssm 2.2.6.post3 (fused add+RMSNorm, gated RMSNorm, Mamba2 SSD
    chunk scan via Hydra's mem-efficient path).
  * ``pure_torch=True``: the same math in plain PyTorch (explicit quadratic SSD
    form, plain RMSNorm). Exportable to ONNX/TensorRT and used as a fallback on
    GPUs where the Triton kernels misbehave. Parity is checked in the tests.
Deviation from the source (deliberate): torchaudio's bundle wrapper applies
``layer_norm(waveforms, waveforms.shape)`` -- jointly over the whole (B, T)
batch, so the source's scores depend on batch composition. We normalise each
utterance on its own (== the source at batch size 1), making scores
batch-independent and reproducible.
"""
from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat

_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)  # vendored mamba_ssm + _hydra resolve from here

from mamba_ssm.modules.mha import MHA  # noqa: E402
from mamba_ssm.ops.triton.layer_norm import RMSNorm, rms_norm_fn  # noqa: E402
from _hydra import Hydra  # noqa: E402

WINDOW = 66800  # source default --fixed-length-samples (4.175 s @ 16 kHz)


@dataclass
class MamBoConfig:
    d_model: int = 128
    d_state: int = 64
    d_conv: int = 7          # xlsr_Hydra task: MODEL_CONFIG.replace(d_conv=7)
    expand: int = 2
    mha_num_heads: int = 4
    mha_mlp_dim: int = 0
    mha_d_conv: int = 4
    residual_in_fp32: bool = True
    use_mem_eff_path: bool = True
    block_layers: int = 5
    n_mamba: int = 3


# ----------------------------------------------------------------------------
# pure-torch equivalents of the Triton ops
# ----------------------------------------------------------------------------
def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    xf = x.float()
    rstd = torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
    return (xf * rstd * weight.float()).to(x.dtype)


def add_rmsnorm(norm: RMSNorm, hidden, residual, prenorm: bool, residual_in_fp32: bool, pure: bool):
    """mamba_ssm fused add + RMSNorm (rms_norm_fn) semantics."""
    if not pure:
        return rms_norm_fn(hidden, norm.weight, norm.bias, residual=residual, prenorm=prenorm,
                           residual_in_fp32=residual_in_fp32, eps=norm.eps)
    res = hidden if residual is None else hidden + residual
    if residual_in_fp32:
        res = res.float()
    out = _rms(res, norm.weight, norm.eps).to(hidden.dtype)
    return (out, res) if prenorm else out


def ssd_scan_torch(x, dt, A, B, C):
    """Mamba2 SSD scan, exact quadratic form (no chunking; L is ~208 here).
    x: (b, l, h, p)  dt: (b, l, h) (already softplus'ed)  A: (h,)
    B, C: (b, l, g, n) with g == 1.  Returns y: (b, l, h, p), fp32 math.
    y_t = sum_{s<=t} (C_t . B_s) * exp(sum_{r=s+1..t} A dt_r) * dt_s * x_s
    """
    dtype = x.dtype
    x, dt, B, C = x.float(), dt.float(), B.float(), C.float()
    dA = dt * A.float()[None, None, :]                       # (b, l, h)
    cs = torch.cumsum(dA, dim=1)                              # (b, l, h)
    seg = cs[:, :, None, :] - cs[:, None, :, :]               # (b, t, s, h)
    L = x.shape[1]
    causal = torch.ones(L, L, dtype=torch.bool, device=x.device).tril()
    decay = torch.exp(seg.masked_fill(~causal[None, :, :, None], float("-inf")))
    CB = torch.einsum("btgn,bsgn->bts", C, B)                 # g == 1
    w = CB[..., None] * decay * dt[:, None, :, :]             # (b, t, s, h)
    y = torch.einsum("btsh,bshp->bthp", w, x)
    return y.to(dtype)


def _chunk_flip_join(t, dim, op):
    h = t.shape[dim] // 2  # static under tracing; slicing (not chunk) -> no ONNX Sequence ops
    fwd, bwd = t.narrow(dim, 0, h), t.narrow(dim, h, h)
    bwd = torch.flip(bwd, dims=(1,))
    if op == "sum":
        return fwd + bwd
    if op == "vstack":
        return torch.cat([fwd, bwd], dim=0)
    raise ValueError(op)


def hydra_forward_torch(m: Hydra, u: torch.Tensor) -> torch.Tensor:
    """Pure-torch mirror of _hydra/ops.py HydraSplitConv1dScanCombinedFn.forward
    (the use_mem_eff_path=True path the checkpoint was trained/evaluated with)."""
    batch, seqlen, _ = u.shape
    d_inner, d_state, headdim, ngroups = m.d_inner, m.d_state, m.headdim, m.ngroups
    nheads = d_inner // headdim
    d_xBC = d_inner + 2 * (2 * ngroups * d_state)
    zxbcdt = m.in_proj(u)
    A = -torch.exp(m.A_log.float())
    z = zxbcdt[..., :d_inner]
    xBC_og = zxbcdt[..., d_inner:d_inner + d_xBC]
    dt = zxbcdt[..., d_inner + d_xBC:]
    k = int(m.d_conv)  # odd (7): "same" padding == k // 2 each side (python int: traceable)
    xBC = F.silu(F.conv1d(xBC_og.transpose(1, 2), m.conv1d.weight, m.conv1d.bias, padding=k // 2,
                          groups=int(m.conv1d.groups)).transpose(1, 2))
    x, BC = xBC[..., :d_inner], xBC[..., d_inner:]
    x_og = x
    x = torch.cat([x, torch.flip(x, dims=(1,))], dim=0)
    BC = _chunk_flip_join(BC, dim=-1, op="vstack")
    nBC = ngroups * d_state
    Bm, Cm = BC[..., :nBC], BC[..., nBC:]
    x = rearrange(x, "b l (h p) -> b l h p", p=headdim)
    Bm = rearrange(Bm, "b l (g n) -> b l g n", g=ngroups)
    Cm = rearrange(Cm, "b l (g n) -> b l g n", g=ngroups)
    D = repeat(F.linear(x_og, m.fc_D.weight, m.D), "b l h -> b l (h p)", p=headdim)
    dt = _chunk_flip_join(dt, dim=-1, op="vstack")
    dt = F.softplus(dt.float() + m.dt_bias.float())
    scan = ssd_scan_torch(x, dt, A, Bm, Cm)
    scan = rearrange(scan, "b l h p -> b l (h p)")
    y = torch.roll(scan, shifts=1, dims=1)
    y = torch.cat([torch.zeros_like(y[:, :1]), y[:, 1:]], dim=1)  # y[:, 0] = 0 (export-friendly)
    y = _chunk_flip_join(y, dim=0, op="sum") + D * x_og
    # gated RMSNorm, norm_before_gate=True
    u_ = _rms(y, m.norm.weight, m.norm.eps).float() * F.silu(z.float())
    return F.linear(u_.to(y.dtype), m.out_proj.weight, m.out_proj.bias)


# ----------------------------------------------------------------------------
# blocks (Add -> RMSNorm -> mixer), names as in step2_model_blocks.py
# ----------------------------------------------------------------------------
class MLP_SwiGLU(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        hidden_dim = int(8 * d_model / 3)
        self.w1_w3_combined = nn.Linear(d_model, 2 * hidden_dim, bias=False)
        self.w2_out = nn.Linear(hidden_dim, d_model, bias=False)
        self.activation = nn.SiLU()

    def forward(self, x):
        gv = self.w1_w3_combined(x)
        h = self.w2_out.in_features
        gate, value = gv[..., :h], gv[..., h:]
        return self.w2_out(self.activation(gate) * value)  # dropout(0.1) is a no-op at eval


class _PreNormBlock(nn.Module):
    pure = False

    def __init__(self, cfg: MamBoConfig):
        super().__init__()
        self.norm = RMSNorm(cfg.d_model)
        self.residual_in_fp32 = cfg.residual_in_fp32

    def mix(self, h):
        raise NotImplementedError

    def forward(self, hidden, residual=None):
        hidden, residual = add_rmsnorm(self.norm, hidden, residual, True, self.residual_in_fp32, self.pure)
        return self.mix(hidden), residual


class Hydra_Block(_PreNormBlock):
    def __init__(self, cfg: MamBoConfig):
        super().__init__(cfg)
        # NB: source passes no headdim -> Hydra default 64 (d_inner 256, 4 heads)
        self.mamba = Hydra(d_model=cfg.d_model, d_state=cfg.d_state, d_conv=cfg.d_conv,
                           expand=cfg.expand, use_mem_eff_path=cfg.use_mem_eff_path)

    def mix(self, h):
        return hydra_forward_torch(self.mamba, h) if self.pure else self.mamba(h)


class MLP_block(_PreNormBlock):
    def __init__(self, cfg: MamBoConfig):
        super().__init__(cfg)
        self.mlp = MLP_SwiGLU(cfg.d_model)

    def mix(self, h):
        return self.mlp(h)


class MHA_Block(_PreNormBlock):
    def __init__(self, cfg: MamBoConfig):
        super().__init__(cfg)
        self.mha = MHA(embed_dim=cfg.d_model, num_heads=cfg.mha_num_heads,
                       mlp_dim=cfg.mha_mlp_dim, d_conv=cfg.mha_d_conv)

    def mix(self, h):
        # MHA without flash_attn/causal_conv1d: nn.Conv1d (causal trim) + SDPA (non-causal)
        return self.mha(h)


class Mix_n_Hydra_1MLP_1MHA_1MLP(nn.Module):
    def __init__(self, cfg: MamBoConfig):
        super().__init__()
        self.layers = nn.ModuleList()
        for _ in range(cfg.block_layers):
            self.layers.extend([Hydra_Block(cfg) for _ in range(cfg.n_mamba)])
            self.layers.append(MLP_block(cfg))
            self.layers.append(MHA_Block(cfg))
            self.layers.append(MLP_block(cfg))
        self.norm = RMSNorm(cfg.d_model)
        self.residual_in_fp32 = cfg.residual_in_fp32
        self.pure = False

    def forward(self, hidden):
        residual = None
        for layer in self.layers:
            hidden, residual = layer(hidden, residual)
        return add_rmsnorm(self.norm, hidden, residual, False, self.residual_in_fp32, self.pure)


class StackedMamba2(nn.Module):
    def __init__(self, cfg: MamBoConfig):
        super().__init__()
        self.layers = Mix_n_Hydra_1MLP_1MHA_1MLP(cfg)

    def forward(self, x):
        return self.layers(x)


class RMS_InputProj(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.input_proj = nn.Linear(1024, d_model, bias=False)
        self.norm = RMSNorm(d_model)
        self.pure = False

    def forward(self, x):
        x = self.input_proj(x)
        x = _rms(x, self.norm.weight, self.norm.eps) if self.pure else self.norm(x)
        return F.silu(x)


class GatedAttentionPool(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        hidden_dim = int(8 * d_model // 3)
        self.gate_value_proj = nn.Linear(d_model, 2 * hidden_dim, bias=False)
        self.to_score = nn.Linear(hidden_dim, 1, bias=False)
        self.activation = nn.SiLU()

    def forward(self, x):
        gv = self.gate_value_proj(x)
        h = self.to_score.in_features
        gate, value = gv[..., :h], gv[..., h:]
        scores = self.to_score(self.activation(gate) * value).squeeze(-1)  # (B, L)
        weights = F.softmax(scores, dim=1).unsqueeze(1)                     # (B, 1, L)
        return torch.bmm(weights, x).squeeze(1)                            # (B, D)


class _BundleWrapper(nn.Module):
    """Key-compatible stand-in for torchaudio's pipelines `_Wav2Vec2Model` wrapper
    (keys `model.*`), but with PER-UTTERANCE waveform normalisation."""

    def __init__(self):
        super().__init__()
        import torchaudio
        # architecture only; the fine-tuned checkpoint overwrites every weight
        self.model = torchaudio.models.wav2vec2_xlsr_300m()

    def forward(self, wav):
        # per-utterance layer_norm (see module docstring), written out explicitly
        # (== F.layer_norm(wav, wav.shape[-1:]), eps 1e-5) so it exports to ONNX.
        mean = wav.mean(dim=-1, keepdim=True)
        var = (wav - mean).square().mean(dim=-1, keepdim=True)
        wav = (wav - mean) * torch.rsqrt(var + 1e-5)
        out, _ = self.model(wav)
        return out


class WLSR300M_Extractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = _BundleWrapper()

    def forward(self, wav):
        return self.model(wav)  # (B, L, 1024), last transformer layer


class XLSRMamBo(nn.Module):
    def __init__(self, cfg: MamBoConfig | None = None, pure_torch: bool = False):
        super().__init__()
        cfg = cfg or MamBoConfig()
        self.feature_extractor = WLSR300M_Extractor()
        self.input_proj = RMS_InputProj(cfg.d_model)
        self.stacked_mamba2 = StackedMamba2(cfg)
        self.attention_pool = GatedAttentionPool(cfg.d_model)
        self.classifier = nn.Linear(cfg.d_model, 2)
        self.set_pure_torch(pure_torch)

    def set_pure_torch(self, flag: bool) -> None:
        for mod in self.modules():
            if hasattr(mod, "pure"):
                mod.pure = flag

    def forward(self, wav: torch.Tensor) -> torch.Tensor:
        x = self.feature_extractor(wav)
        x = self.input_proj(x)
        x = self.stacked_mamba2(x)
        x = self.attention_pool(x)
        return self.classifier(x)  # (B, 2) logits; class 1 = bona fide


def load_state(path: str) -> dict:
    sd = torch.load(path, map_location="cpu", weights_only=True)
    return {(k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v for k, v in sd.items()}


def pad_repeat(wav, max_len: int = WINDOW):
    """Source pad_wave_repeat: first max_len samples; shorter clips are tiled."""
    import numpy as np
    n = len(wav)
    if n == 0:
        raise ValueError("empty waveform")
    if n >= max_len:
        return wav[:max_len]
    return np.tile(wav, -(-max_len // n))[:max_len]

	"""XLSR-MamBo network (MamBo-3: Hydra-N3, 5 repeats), ported from
	https://github.com/saki-ciallo/XLSR-MamBo (MIT) @ f429bb1:

	- model_build.F4_XLSR300M_General -> XLSRMamBo
	- modules/step1_feature_ptm.WLSR300M_Extractor + RMS_InputProj
	- modules/step2_model_blocks.{Hydra_Block, MLP_block, MHA_Block, MLP_SwiGLU,
	Mix_n_Hydra_1MLP_1MHA_1MLP}
	- modules/step3_pooling_outputs.GatedAttentionPool
	- modules/hydra (Hydra mixer, vendored in ./_hydra)
	- train_config: task xlsr_Hydra = Mamba2Config(d_model=128, d_state=64, d_conv=7,
	expand=2, mha_num_heads=4, mha_mlp_dim=0, mha_d_conv=4, residual_in_fp32=True,
	fused_add_norm=True, use_mem_eff_path=True); block_layers=5, n_mamba=3.

	Parameter/module names are kept identical to the source so the released
	state_dict loads strictly.

	Two numerically-equivalent backends:
	* ``pure_torch=False`` (default): the source path -- Triton kernels from the
	vendored mamba_ssm 2.2.6.post3 (fused add+RMSNorm, gated RMSNorm, Mamba2 SSD
	chunk scan via Hydra's mem-efficient path).
	* ``pure_torch=True``: the same math in plain PyTorch (explicit quadratic SSD
	form, plain RMSNorm). Exportable to ONNX/TensorRT and used as a fallback on
	GPUs where the Triton kernels misbehave. Parity is checked in the tests.

	Deviation from the source (deliberate): torchaudio's bundle wrapper applies
	``layer_norm(waveforms, waveforms.shape)`` -- jointly over the whole (B, T)
	batch, so the source's scores depend on batch composition. We normalise each
	utterance on its own (== the source at batch size 1), making scores
	batch-independent and reproducible.
	"""
	from __future__ import annotations

	import os
	import sys
	from dataclasses import dataclass
	from typing import Optional, Tuple

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from einops import rearrange, repeat

	_HERE = os.path.dirname(os.path.abspath(__file__))
	if _HERE not in sys.path:
	sys.path.insert(0, _HERE) # vendored mamba_ssm + _hydra resolve from here

	from mamba_ssm.modules.mha import MHA # noqa: E402
	from mamba_ssm.ops.triton.layer_norm import RMSNorm, rms_norm_fn # noqa: E402
	from _hydra import Hydra # noqa: E402

	WINDOW = 66800 # source default --fixed-length-samples (4.175 s @ 16 kHz)


	@dataclass
	class MamBoConfig:
	d_model: int = 128
	d_state: int = 64
	d_conv: int = 7 # xlsr_Hydra task: MODEL_CONFIG.replace(d_conv=7)
	expand: int = 2
	mha_num_heads: int = 4
	mha_mlp_dim: int = 0
	mha_d_conv: int = 4
	residual_in_fp32: bool = True
	use_mem_eff_path: bool = True
	block_layers: int = 5
	n_mamba: int = 3


	# ----------------------------------------------------------------------------
	# pure-torch equivalents of the Triton ops
	# ----------------------------------------------------------------------------
	def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
	xf = x.float()
	rstd = torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
	return (xf * rstd * weight.float()).to(x.dtype)


	def add_rmsnorm(norm: RMSNorm, hidden, residual, prenorm: bool, residual_in_fp32: bool, pure: bool):
	"""mamba_ssm fused add + RMSNorm (rms_norm_fn) semantics."""
	if not pure:
	return rms_norm_fn(hidden, norm.weight, norm.bias, residual=residual, prenorm=prenorm,
	residual_in_fp32=residual_in_fp32, eps=norm.eps)
	res = hidden if residual is None else hidden + residual
	if residual_in_fp32:
	res = res.float()
	out = _rms(res, norm.weight, norm.eps).to(hidden.dtype)
	return (out, res) if prenorm else out


	def ssd_scan_torch(x, dt, A, B, C):
	"""Mamba2 SSD scan, exact quadratic form (no chunking; L is ~208 here).

	x: (b, l, h, p) dt: (b, l, h) (already softplus'ed) A: (h,)
	B, C: (b, l, g, n) with g == 1. Returns y: (b, l, h, p), fp32 math.
	y_t = sum_{s<=t} (C_t . B_s) * exp(sum_{r=s+1..t} A dt_r) * dt_s * x_s
	"""
	dtype = x.dtype
	x, dt, B, C = x.float(), dt.float(), B.float(), C.float()
	dA = dt * A.float()[None, None, :] # (b, l, h)
	cs = torch.cumsum(dA, dim=1) # (b, l, h)
	seg = cs[:, :, None, :] - cs[:, None, :, :] # (b, t, s, h)
	L = x.shape[1]
	causal = torch.ones(L, L, dtype=torch.bool, device=x.device).tril()
	decay = torch.exp(seg.masked_fill(~causal[None, :, :, None], float("-inf")))
	CB = torch.einsum("btgn,bsgn->bts", C, B) # g == 1
	w = CB[..., None] * decay * dt[:, None, :, :] # (b, t, s, h)
	y = torch.einsum("btsh,bshp->bthp", w, x)
	return y.to(dtype)


	def _chunk_flip_join(t, dim, op):
	h = t.shape[dim] // 2 # static under tracing; slicing (not chunk) -> no ONNX Sequence ops
	fwd, bwd = t.narrow(dim, 0, h), t.narrow(dim, h, h)
	bwd = torch.flip(bwd, dims=(1,))
	if op == "sum":
	return fwd + bwd
	if op == "vstack":
	return torch.cat([fwd, bwd], dim=0)
	raise ValueError(op)


	def hydra_forward_torch(m: Hydra, u: torch.Tensor) -> torch.Tensor:
	"""Pure-torch mirror of _hydra/ops.py HydraSplitConv1dScanCombinedFn.forward
	(the use_mem_eff_path=True path the checkpoint was trained/evaluated with)."""
	batch, seqlen, _ = u.shape
	d_inner, d_state, headdim, ngroups = m.d_inner, m.d_state, m.headdim, m.ngroups
	nheads = d_inner // headdim
	d_xBC = d_inner + 2 * (2 * ngroups * d_state)
	zxbcdt = m.in_proj(u)
	A = -torch.exp(m.A_log.float())
	z = zxbcdt[..., :d_inner]
	xBC_og = zxbcdt[..., d_inner:d_inner + d_xBC]
	dt = zxbcdt[..., d_inner + d_xBC:]
	k = int(m.d_conv) # odd (7): "same" padding == k // 2 each side (python int: traceable)
	xBC = F.silu(F.conv1d(xBC_og.transpose(1, 2), m.conv1d.weight, m.conv1d.bias, padding=k // 2,
	groups=int(m.conv1d.groups)).transpose(1, 2))
	x, BC = xBC[..., :d_inner], xBC[..., d_inner:]
	x_og = x
	x = torch.cat([x, torch.flip(x, dims=(1,))], dim=0)
	BC = _chunk_flip_join(BC, dim=-1, op="vstack")
	nBC = ngroups * d_state
	Bm, Cm = BC[..., :nBC], BC[..., nBC:]
	x = rearrange(x, "b l (h p) -> b l h p", p=headdim)
	Bm = rearrange(Bm, "b l (g n) -> b l g n", g=ngroups)
	Cm = rearrange(Cm, "b l (g n) -> b l g n", g=ngroups)
	D = repeat(F.linear(x_og, m.fc_D.weight, m.D), "b l h -> b l (h p)", p=headdim)
	dt = _chunk_flip_join(dt, dim=-1, op="vstack")
	dt = F.softplus(dt.float() + m.dt_bias.float())
	scan = ssd_scan_torch(x, dt, A, Bm, Cm)
	scan = rearrange(scan, "b l h p -> b l (h p)")
	y = torch.roll(scan, shifts=1, dims=1)
	y = torch.cat([torch.zeros_like(y[:, :1]), y[:, 1:]], dim=1) # y[:, 0] = 0 (export-friendly)
	y = _chunk_flip_join(y, dim=0, op="sum") + D * x_og
	# gated RMSNorm, norm_before_gate=True
	u_ = _rms(y, m.norm.weight, m.norm.eps).float() * F.silu(z.float())
	return F.linear(u_.to(y.dtype), m.out_proj.weight, m.out_proj.bias)


	# ----------------------------------------------------------------------------
	# blocks (Add -> RMSNorm -> mixer), names as in step2_model_blocks.py
	# ----------------------------------------------------------------------------
	class MLP_SwiGLU(nn.Module):
	def __init__(self, d_model: int):
	super().__init__()
	hidden_dim = int(8 * d_model / 3)
	self.w1_w3_combined = nn.Linear(d_model, 2 * hidden_dim, bias=False)
	self.w2_out = nn.Linear(hidden_dim, d_model, bias=False)
	self.activation = nn.SiLU()

	def forward(self, x):
	gv = self.w1_w3_combined(x)
	h = self.w2_out.in_features
	gate, value = gv[..., :h], gv[..., h:]
	return self.w2_out(self.activation(gate) * value) # dropout(0.1) is a no-op at eval


	class _PreNormBlock(nn.Module):
	pure = False

	def __init__(self, cfg: MamBoConfig):
	super().__init__()
	self.norm = RMSNorm(cfg.d_model)
	self.residual_in_fp32 = cfg.residual_in_fp32

	def mix(self, h):
	raise NotImplementedError

	def forward(self, hidden, residual=None):
	hidden, residual = add_rmsnorm(self.norm, hidden, residual, True, self.residual_in_fp32, self.pure)
	return self.mix(hidden), residual


	class Hydra_Block(_PreNormBlock):
	def __init__(self, cfg: MamBoConfig):
	super().__init__(cfg)
	# NB: source passes no headdim -> Hydra default 64 (d_inner 256, 4 heads)
	self.mamba = Hydra(d_model=cfg.d_model, d_state=cfg.d_state, d_conv=cfg.d_conv,
	expand=cfg.expand, use_mem_eff_path=cfg.use_mem_eff_path)

	def mix(self, h):
	return hydra_forward_torch(self.mamba, h) if self.pure else self.mamba(h)


	class MLP_block(_PreNormBlock):
	def __init__(self, cfg: MamBoConfig):
	super().__init__(cfg)
	self.mlp = MLP_SwiGLU(cfg.d_model)

	def mix(self, h):
	return self.mlp(h)


	class MHA_Block(_PreNormBlock):
	def __init__(self, cfg: MamBoConfig):
	super().__init__(cfg)
	self.mha = MHA(embed_dim=cfg.d_model, num_heads=cfg.mha_num_heads,
	mlp_dim=cfg.mha_mlp_dim, d_conv=cfg.mha_d_conv)

	def mix(self, h):
	# MHA without flash_attn/causal_conv1d: nn.Conv1d (causal trim) + SDPA (non-causal)
	return self.mha(h)


	class Mix_n_Hydra_1MLP_1MHA_1MLP(nn.Module):
	def __init__(self, cfg: MamBoConfig):
	super().__init__()
	self.layers = nn.ModuleList()
	for _ in range(cfg.block_layers):
	self.layers.extend([Hydra_Block(cfg) for _ in range(cfg.n_mamba)])
	self.layers.append(MLP_block(cfg))
	self.layers.append(MHA_Block(cfg))
	self.layers.append(MLP_block(cfg))
	self.norm = RMSNorm(cfg.d_model)
	self.residual_in_fp32 = cfg.residual_in_fp32
	self.pure = False

	def forward(self, hidden):
	residual = None
	for layer in self.layers:
	hidden, residual = layer(hidden, residual)
	return add_rmsnorm(self.norm, hidden, residual, False, self.residual_in_fp32, self.pure)


	class StackedMamba2(nn.Module):
	def __init__(self, cfg: MamBoConfig):
	super().__init__()
	self.layers = Mix_n_Hydra_1MLP_1MHA_1MLP(cfg)

	def forward(self, x):
	return self.layers(x)


	class RMS_InputProj(nn.Module):
	def __init__(self, d_model: int):
	super().__init__()
	self.input_proj = nn.Linear(1024, d_model, bias=False)
	self.norm = RMSNorm(d_model)
	self.pure = False

	def forward(self, x):
	x = self.input_proj(x)
	x = _rms(x, self.norm.weight, self.norm.eps) if self.pure else self.norm(x)
	return F.silu(x)


	class GatedAttentionPool(nn.Module):
	def __init__(self, d_model: int):
	super().__init__()
	hidden_dim = int(8 * d_model // 3)
	self.gate_value_proj = nn.Linear(d_model, 2 * hidden_dim, bias=False)
	self.to_score = nn.Linear(hidden_dim, 1, bias=False)
	self.activation = nn.SiLU()

	def forward(self, x):
	gv = self.gate_value_proj(x)
	h = self.to_score.in_features
	gate, value = gv[..., :h], gv[..., h:]
	scores = self.to_score(self.activation(gate) * value).squeeze(-1) # (B, L)
	weights = F.softmax(scores, dim=1).unsqueeze(1) # (B, 1, L)
	return torch.bmm(weights, x).squeeze(1) # (B, D)


	class _BundleWrapper(nn.Module):
	"""Key-compatible stand-in for torchaudio's pipelines `_Wav2Vec2Model` wrapper
	(keys `model.*`), but with PER-UTTERANCE waveform normalisation."""

	def __init__(self):
	super().__init__()
	import torchaudio
	# architecture only; the fine-tuned checkpoint overwrites every weight
	self.model = torchaudio.models.wav2vec2_xlsr_300m()

	def forward(self, wav):
	# per-utterance layer_norm (see module docstring), written out explicitly
	# (== F.layer_norm(wav, wav.shape[-1:]), eps 1e-5) so it exports to ONNX.
	mean = wav.mean(dim=-1, keepdim=True)
	var = (wav - mean).square().mean(dim=-1, keepdim=True)
	wav = (wav - mean) * torch.rsqrt(var + 1e-5)
	out, _ = self.model(wav)
	return out


	class WLSR300M_Extractor(nn.Module):
	def __init__(self):
	super().__init__()
	self.model = _BundleWrapper()

	def forward(self, wav):
	return self.model(wav) # (B, L, 1024), last transformer layer


	class XLSRMamBo(nn.Module):
	def __init__(self, cfg: MamBoConfig \| None = None, pure_torch: bool = False):
	super().__init__()
	cfg = cfg or MamBoConfig()
	self.feature_extractor = WLSR300M_Extractor()
	self.input_proj = RMS_InputProj(cfg.d_model)
	self.stacked_mamba2 = StackedMamba2(cfg)
	self.attention_pool = GatedAttentionPool(cfg.d_model)
	self.classifier = nn.Linear(cfg.d_model, 2)
	self.set_pure_torch(pure_torch)

	def set_pure_torch(self, flag: bool) -> None:
	for mod in self.modules():
	if hasattr(mod, "pure"):
	mod.pure = flag

	def forward(self, wav: torch.Tensor) -> torch.Tensor:
	x = self.feature_extractor(wav)
	x = self.input_proj(x)
	x = self.stacked_mamba2(x)
	x = self.attention_pool(x)
	return self.classifier(x) # (B, 2) logits; class 1 = bona fide


	def load_state(path: str) -> dict:
	sd = torch.load(path, map_location="cpu", weights_only=True)
	return {(k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v for k, v in sd.items()}


	def pad_repeat(wav, max_len: int = WINDOW):
	"""Source pad_wave_repeat: first max_len samples; shorter clips are tiled."""
	import numpy as np
	n = len(wav)
	if n == 0:
	raise ValueError("empty waveform")
	if n >= max_len:
	return wav[:max_len]
	return np.tile(wav, -(-max_len // n))[:max_len]