remote-code bundle

modeling_qwen3_mamba3.py

@@ -0,0 +1,1002 @@
+# GENERATED FILE -- do not edit.
+#
+# Produced by `scripts/mamba3/make_hf_remote_code_bundle.py` from
+# `src/models/qwen3_mamba3.py` in the Linearization repo (branch mamba3). Edit that
+# file and regenerate. Carries the config and every model class; config.json's
+# `auto_map` points AutoConfig and AutoModelForCausalLM here.
+#
+# The fused Mamba-3 MIMO kernel (state-spaces/mamba at e9594ce plus the local
+# backward fix described in the model card) is optional: without it every mixer
+# runs `mamba3_mimo_reference`, a sequential fp32 torch path that is exact but slow.
+
+"""Qwen3 -> Mamba-3 (MIMO) hybrid container.
+
+The Mamba-3 sibling of ``src/models/qwen3_mamba2.py``: retained layers are literal
+``Qwen3DecoderLayer``s (bit-exact retention, no masquerade), converted layers mount a
+Mamba-3 MIMO mixer under the load-bearing attribute name ``linear_attn`` (all four
+trainers split LR groups on the ``".linear_attn."`` substring, and
+``stage1_align.py::layer_forward`` calls ``layer.linear_attn(hidden_states=...)`` and
+expects a bare tensor back).
+
+Why a hand-written mixer rather than the upstream ``mamba_ssm.modules.mamba3.Mamba3``:
+
+* upstream's module is inseparable from its kernels -- it asserts the TileLang MIMO
+  kernel imported at construction and has no torch fallback -- so nothing built on it
+  could run the CPU surgery tests, the login-node surgery, or a cached decode step
+  (its ``step()`` needs the CuteDSL ``mamba3_step_fn``, which needs ``quack``, which
+  this env does not have);
+* its ``_preprocess`` for decode does ``B.expand(-1, -1, nheads, -1)``, which is only
+  valid for ``ngroups == 1``; we run 8 B/C groups (one per teacher KV head);
+* every init choice here is deliberate (see ``Qwen3Mamba3Mixer.init_ssm_parameters``)
+  and several differ from upstream's -- the whole point of the surgery is that the
+  mixer starts as the teacher's attention head, not as a random SSM.
+
+The **math and the parameter layout are upstream's exactly** (mamba_ssm @ e9594ce,
+``mamba_ssm/modules/mamba3.py``): same ``in_proj`` row order
+``[z | x | B | C | dd_dt | dd_A | trap | angles]``, same rank-major ``(r g n)`` B/C
+layout, same ``dt_bias``/``B_bias``/``C_bias``/``mimo_x``/``mimo_z``/``mimo_o``/``D``/
+``norm.weight`` parameter names and shapes, and the prefill calls upstream's fused
+TileLang kernel ``mamba3_mimo`` untouched. The pure-torch path
+(``mamba3_mimo_reference``) is a transcription of the kernel's own semantics -- the
+rotate-half pairing ``(n, N/2 + n)`` for ``n < N // rotary_dim_divisor``, bias added
+before rotation, exponential-trapezoidal state update, per-(head, rank) pre-gate
+RMSNorm, ``mimo_o`` reduction -- and is checked against the kernel on a GPU in
+``tests/test_qwen3_mamba3_surgery.py``.
+
+Importing this module is CPU-safe: the kernel import is lazy and only attempted for
+CUDA tensors.
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.bamba.modeling_bamba import HybridMambaAttentionDynamicCache
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3DecoderLayer,
+    Qwen3MLP,
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+)
+
+# Pinned upstream: github.com/state-spaces/mamba @ e9594ce1c732d97440f0332fdc43170a2294dbfa
+# plus one local commit in the staged clone, cc12e795314d31a2e68e1407bd4309a91673bc23:
+# mamba3_mimo backward reads ``ctx.saved_tensors`` once, so it survives HF non-reentrant
+# activation checkpointing (stage 2 KD). Stage 1 never checkpoints, so it never hit this.
+# (staged at /mnt/home/afei1/mamba3_upstream; scripts/beta/_env_b200_mamba3.sh prepends
+# it to PYTHONPATH so the whole ``mamba_ssm`` package resolves there).
+UPSTREAM_SHA = "cc12e795314d31a2e68e1407bd4309a91673bc23"  # e9594ce + local checkpoint-compat fix
+
+
+class Qwen3Mamba3Config(Qwen3Config):
+    """Qwen3 config + the Mamba-3 MIMO parameterization + hybrid ``layer_types``.
+
+    The Qwen3-4B backbone is restated as defaults for the same reason as in the
+    Mamba2 container: ``PretrainedConfig.to_diff_dict`` builds a zero-arg instance to
+    diff against, and the consistency checks below must hold for that instance or
+    ``save_pretrained`` raises.
+
+    Geometry (Qwen3-4B, hidden 2560, 32 q heads x 128, 8 kv heads):
+
+    * ``mamba_n_heads`` = 32 = teacher q heads; ``mamba_d_head`` = 128 (P, the value
+      width); ``mamba_d_state`` = 128 (N, the key width -- scores are C . B).
+    * ``mamba_n_groups`` = 8 = teacher kv heads. Head h reads B/C group ``h // (H/G)``
+      (``mamba3_mimo_fwd.py:142``, ``i_h_qk = i_h // (H // G)``), which is exactly
+      ``repeat_kv``: the four heads of a group share one key, like GQA.
+    * ``mamba_mimo_rank`` = 4 = gqa. Rank r of group g carries teacher q head
+      ``4g + r``: the four q heads that share a kv head become the four MIMO ranks
+      of that group, so the MIMO structure is *populated by GQA* rather than by
+      random init. Any R with ``R | gqa`` or ``gqa | R`` is legal (heads are
+      averaged or replicated across ranks); R=2 is the parameter-matched control.
+    * ``mamba_rope_fraction`` in {0.5, 1.0}: fraction of N carried through the
+      data-dependent rotation (kernel supports divisors 4 and 2 only). Only 0.5 is
+      tested upstream.
+    * ``mamba_chunk_size``: kernel chunk. The kernels tile ``fused = chunk * R`` rows and
+      the backward-of-backward kernel holds every tile in shared memory at once
+      (:func:`mamba3_bwd_smem_bytes`); at N=P=128 that overflows the 227 KiB a B200/H100
+      block may use as soon as ``fused > 32`` (chunk 16 x R=4 asked for 281248 B and did
+      not compile in the first stage-1 gate, job 196032). From below, ``fused < 32`` does
+      not compile either: the kernels' ``fused x fused`` GEMMs must split across the
+      block's 4 warps and TileLang's MMA warp partition cannot tile 16 x 16 four ways
+      (chunk 8 x R=2, job 196045). So at N=P=128 ``chunk * R == 32`` exactly
+      (:data:`MAMBA3_MIN_FUSED_TILE`): default 8 for R=4; R=2 uses 16
+      (:func:`mamba3_max_chunk_size`). Upstream's ``64 / R`` advice assumes its headdim
+      64. Tiling only: the math is chunk-invariant. The kernel asserts ``>= 8``.
+    * ``mamba_bc_bias_init``: value the B/C biases start at. Upstream uses 1.0. The
+      worry that on an inherited init this adds ``sum(q) + sum(k) + N`` to every
+      attention score and swamps the teacher's structure was tested (bias 0 vs 1, same
+      surgery): bias 1 led at every stage-1 and stage-2a eval (0.00376 vs 0.00397 rel-MSE,
+      0.1258 vs 0.1297 KL), so the default is upstream's 1.0.
+    """
+
+    model_type = "qwen3_mamba3"
+
+    def __init__(
+        self,
+        # --- Qwen3-4B backbone, restated so the zero-arg constructor is valid ---
+        vocab_size: int = 151936,
+        hidden_size: int = 2560,
+        intermediate_size: int = 9728,
+        num_hidden_layers: int = 36,
+        num_attention_heads: int = 32,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        # --- Mamba-3 MIMO parameterization ---
+        mamba_n_heads: int = 32,
+        mamba_d_head: int = 128,
+        mamba_d_state: int = 128,
+        mamba_n_groups: int = 8,
+        mamba_mimo_rank: int = 4,
+        mamba_expand: float = 1.6,        # int(1.6 * 2560) == 4096 exactly
+        mamba_chunk_size: int = 8,
+        mamba_rope_fraction: float = 0.5,
+        mamba_bc_bias_init: float = 1.0,
+        mamba_a_floor: float = 1e-4,
+        mamba_dt_min: float = 0.001,
+        mamba_dt_max: float = 0.1,
+        mamba_dt_init_floor: float = 1e-4,
+        mamba_norm_eps: float = 1e-5,       # pre-gate per-head RMSNorm (upstream 1e-5)
+        mamba_bc_norm_eps: Optional[float] = None,  # None -> rms_norm_eps (teacher q/k_norm)
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            **kwargs,
+        )
+        self.mamba_n_heads = mamba_n_heads
+        self.mamba_d_head = mamba_d_head
+        self.mamba_d_state = mamba_d_state
+        self.mamba_n_groups = mamba_n_groups
+        self.mamba_mimo_rank = mamba_mimo_rank
+        self.mamba_expand = mamba_expand
+        self.mamba_chunk_size = mamba_chunk_size
+        self.mamba_rope_fraction = mamba_rope_fraction
+        self.mamba_bc_bias_init = mamba_bc_bias_init
+        self.mamba_a_floor = mamba_a_floor
+        self.mamba_dt_min = mamba_dt_min
+        self.mamba_dt_max = mamba_dt_max
+        self.mamba_dt_init_floor = mamba_dt_init_floor
+        self.mamba_norm_eps = mamba_norm_eps
+        self.mamba_bc_norm_eps = mamba_bc_norm_eps
+
+        got = int(self.mamba_expand * self.hidden_size)
+        want = self.mamba_n_heads * self.mamba_d_head
+        if got != want:
+            raise ValueError(
+                f"int(mamba_expand * hidden_size) = int({self.mamba_expand} * "
+                f"{self.hidden_size}) = {got}, but mamba_n_heads * mamba_d_head = "
+                f"{self.mamba_n_heads} * {self.mamba_d_head} = {want}. These must be "
+                "equal for out_proj to inherit the teacher's o_proj verbatim."
+            )
+        if self.mamba_n_heads % self.mamba_n_groups != 0:
+            raise ValueError(
+                f"mamba_n_heads ({self.mamba_n_heads}) must be divisible by "
+                f"mamba_n_groups ({self.mamba_n_groups}): head h reads B/C group "
+                "h // (H/G), exactly the way repeat_kv shares GQA kv heads."
+            )
+        hpg = self.mamba_n_heads // self.mamba_n_groups
+        R = self.mamba_mimo_rank
+        if R < 1 or not (hpg % R == 0 or R % hpg == 0):
+            raise ValueError(
+                f"mamba_mimo_rank ({R}) must divide or be a multiple of heads-per-group "
+                f"({hpg}) so ranks map onto the teacher's q heads of a kv group"
+            )
+        if self.mamba_rope_fraction not in (0.5, 1.0):
+            raise ValueError("mamba_rope_fraction must be 0.5 or 1.0 (kernel divisors 4/2)")
+        if self.mamba_d_state % 2 != 0 or mamba3_num_rope_angles(self) < 1:
+            raise ValueError(f"mamba_d_state ({self.mamba_d_state}) must be even and >= 4")
+        if self.mamba_chunk_size < 8:
+            raise ValueError("mamba_chunk_size must be >= 8 (kernel assertion)")
+        fused = self.mamba_chunk_size * self.mamba_mimo_rank
+        if fused < MAMBA3_MIN_FUSED_TILE:
+            raise ValueError(
+                f"mamba_chunk_size={self.mamba_chunk_size} x mamba_mimo_rank={self.mamba_mimo_rank} = "
+                f"fused tile {fused} < {MAMBA3_MIN_FUSED_TILE}: the kernels' fused x fused GEMMs cannot be "
+                f"split across a block's 4 warps (TileLang warp partition; job 196045). Use chunk >= "
+                f"{-(-MAMBA3_MIN_FUSED_TILE // self.mamba_mimo_rank)}"
+            )
+        div = mamba3_rotary_dim_divisor(self)
+        need = mamba3_bwd_smem_bytes(self.mamba_d_state, self.mamba_d_head, self.mamba_chunk_size,
+                                     self.mamba_mimo_rank, div)
+        if need > MAMBA3_MAX_DYN_SMEM_BYTES:
+            fit = mamba3_max_chunk_size(self.mamba_d_state, self.mamba_d_head, self.mamba_mimo_rank, div)
+            raise ValueError(
+                f"mamba_chunk_size={self.mamba_chunk_size} x mamba_mimo_rank={self.mamba_mimo_rank} needs "
+                f"{need} B of shared memory in the MIMO backward kernel at N={self.mamba_d_state}, "
+                f"P={self.mamba_d_head}; a block may use {MAMBA3_MAX_DYN_SMEM_BYTES} B (227 KiB). "
+                f"Largest chunk that fits: {fit}"
+            )
+
+    @property
+    def layers_block_type(self) -> list[str]:
+        """Bamba's vocabulary for ``layer_types`` ("mamba"/"attention")."""
+        return ["mamba" if t == "linear_attention" else "attention" for t in self.layer_types]
+
+
+# Largest dynamic shared memory one block may opt into on sm_90 / sm_100 (227 KiB).
+MAMBA3_MAX_DYN_SMEM_BYTES = 227 * 1024
+# Smallest fused tile (chunk * R) the kernels compile at: every kernel runs 128 threads =
+# 4 warps and TileLang's default MMA warp partition of a fused x fused GEMM needs
+# m_warp * n_warp == 4, which a 16 x 16 tile cannot give (1 x 1; "m_warp * n_warp must
+# equal num_warps", job 196045 at chunk 8 x R=2). 32 x 32 splits 2 x 2 and is what both
+# production geometries run.
+MAMBA3_MIN_FUSED_TILE = 32
+
+
+def mamba3_bwd_smem_bytes(d_state: int, d_head: int, chunk_size: int, mimo_rank: int,
+                          rotary_dim_divisor: int = 4) -> int:
+    """Shared memory (bytes) upstream's ``mamba_mimo_bwd_bwd`` TileLang kernel allocates per block.
+
+    Mirrors its ``alloc_shared`` list (mamba3_mimo_bwd.py; ``fused = chunk * R``; bf16 tiles
+    unless noted): states, dstates [N,P]; dout, v [chunk,P]; dPhiO, PsiV, dPsiV_combined
+    [fused,P]; q, k, q_pre_rot, k_pre_rot, dq, dk, k_pre_trap [fused,N]; lkq_masked and its
+    dtype buffer [fused,fused]; dqk_from_diag [fused,fused] fp32; dangle [fused,N/div] fp32;
+    qk_dot [chunk,R,R]; trap_scale [chunk]; dA_cs, dA_cs_rev [chunk] fp32; plus 2048 B the
+    compiler adds (modeled 279200 vs the 281248 the driver refused at N=P=128, chunk 16,
+    R=4). Upstream pipelines with ``num_stages=0``, so there is no multiplier.
+    """
+    N, P, c, R = d_state, d_head, chunk_size, mimo_rank
+    f = c * R
+    bf, f32 = 2, 4
+    b = 2 * N * P * bf                                  # states, dstates
+    b += 2 * c * P * bf                                 # dout, v
+    b += 3 * f * P * bf                                 # dPhiO, PsiV, dPsiV_combined
+    b += 7 * f * N * bf                                 # q, k, q/k_pre_rot, dq, dk, k_pre_trap
+    b += 2 * f * f * bf + f * f * f32                   # lkq_masked (+ dtype buf), dqk_from_diag
+    b += f * (N // rotary_dim_divisor) * f32            # dangle
+    b += c * R * R * bf + c * bf + 2 * c * f32          # qk_dot, trap_scale, dA_cs, dA_cs_rev
+    return b + 2048
+
+
+def mamba3_max_chunk_size(d_state: int, d_head: int, mimo_rank: int, rotary_dim_divisor: int = 4,
+                          headroom: int = 8 * 1024) -> int | None:
+    """Largest power-of-two chunk in [8, 256] whose backward kernel fits with ``headroom`` spare
+    and whose fused tile ``chunk * R`` reaches :data:`MAMBA3_MIN_FUSED_TILE`.
+
+    N=P=128: 8 for R=4, 16 for R=2, 32 for R=1, None for R=8 (fused 64 overflows at any
+    chunk >= 8).
+    """
+    best, c = None, 8
+    while c <= 256:
+        if c * mimo_rank >= MAMBA3_MIN_FUSED_TILE:
+            if mamba3_bwd_smem_bytes(d_state, d_head, c, mimo_rank, rotary_dim_divisor) + headroom > MAMBA3_MAX_DYN_SMEM_BYTES:
+                break
+            best = c
+        c *= 2
+    return best
+
+
+def mamba3_rotary_dim_divisor(config) -> int:
+    return int(round(2 / config.mamba_rope_fraction))
+
+
+def mamba3_num_rope_angles(config) -> int:
+    """Upstream: ``split = int(d_state * rope_fraction)`` (made even), angles = split // 2."""
+    split = int(config.mamba_d_state * config.mamba_rope_fraction)
+    if split % 2:
+        split -= 1
+    return split // 2
+
+
+def rank_head_offsets(heads_per_group: int, mimo_rank: int) -> list[list[int]]:
+    """For each MIMO rank, the teacher q-head offsets (within one kv group) it carries.
+
+    R == hpg: rank r <-> q head offset r (the production case, 4 == 4).
+    R <  hpg: rank r averages offsets ``[r*m, (r+1)*m)`` with m = hpg // R.
+    R >  hpg: offset j is replicated onto ranks ``[j*m, (j+1)*m)`` with m = R // hpg.
+    Shared by the surgery (which writes the C rows) and ``mimo_o`` init (which reads
+    them back out), so the two cannot disagree.
+    """
+    hpg, R = heads_per_group, mimo_rank
+    if R <= hpg:
+        if hpg % R:
+            raise ValueError(f"mimo_rank {R} does not divide heads-per-group {hpg}")
+        m = hpg // R
+        return [list(range(r * m, (r + 1) * m)) for r in range(R)]
+    if R % hpg:
+        raise ValueError(f"mimo_rank {R} is not a multiple of heads-per-group {hpg}")
+    m = R // hpg
+    return [[r // m] for r in range(R)]
+
+
+def own_rank_matrix(num_heads: int, n_groups: int, mimo_rank: int) -> torch.Tensor:
+    """``(H, R)`` read-out weights making head h output *its own* teacher head.
+
+    Head h = g*hpg + j is teacher attention head h; its query lives on the rank(s)
+    whose offsets include j. Entry ``[h, r]`` is ``1/len(own ranks)`` on those ranks
+    and 0 elsewhere, so at init ``mimo_o`` selects the inherited head instead of
+    averaging the group's four q heads into every head (upstream's ``ones / R``).
+    """
+    hpg = num_heads // n_groups
+    offsets = rank_head_offsets(hpg, mimo_rank)
+    mo = torch.zeros(num_heads, mimo_rank)
+    for h in range(num_heads):
+        j = h % hpg
+        own = [r for r in range(mimo_rank) if j in offsets[r]]
+        mo[h, own] = 1.0 / len(own)
+    return mo
+
+
+def heavy_tail_activation(x: torch.Tensor) -> torch.Tensor:
+    """Upstream ``heavy_tail_activation``: identity for x >= 0, ``1 / (1 - x)`` for x < 0.
+
+    Continuous at 0 with value 1 and slope 1 on both sides, so zero-initialized
+    ``dd_A`` rows give A = -1 for every head and a live gradient in both directions.
+    """
+    return x.clamp(min=0) + 1.0 / (1.0 - x.clamp(max=0))
+
+
+def _rotate_half_partial(t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Kernel-aligned rotation on the last dim: for n < A, pair ``(n, N/2 + n)``.
+
+    ``t``: (..., N); ``cos``/``sin``: broadcastable to (..., A) with A <= N/2.
+    Dims outside the two rotated slices pass through unchanged (rope_fraction 0.5).
+    """
+    N = t.shape[-1]
+    A = cos.shape[-1]
+    half = N // 2
+    t0, mid, t1, rest = t[..., :A], t[..., A:half], t[..., half : half + A], t[..., half + A :]
+    r0 = t0 * cos - t1 * sin
+    r1 = t0 * sin + t1 * cos
+    return torch.cat([r0, mid, r1, rest], dim=-1)
+
+
+def mamba3_mimo_reference(
+    C: torch.Tensor,
+    B: torch.Tensor,
+    x: torch.Tensor,
+    z: torch.Tensor,
+    adt: torch.Tensor,
+    dt: torch.Tensor,
+    trap: torch.Tensor,
+    C_bias: torch.Tensor,
+    B_bias: torch.Tensor,
+    mimo_x: torch.Tensor,
+    mimo_z: torch.Tensor,
+    mimo_o: torch.Tensor,
+    angles: torch.Tensor,
+    D: torch.Tensor,
+    norm_weight: torch.Tensor,
+    norm_eps: float,
+    initial_states: Optional[tuple] = None,
+) -> tuple[torch.Tensor, tuple]:
+    """Sequential fp32 transcription of the fused kernel (prefill *and* decode).
+
+    Shapes (b batch, l tokens, R ranks, G groups, H heads, N state, P head width,
+    A rope angles):
+      C, B: (b, l, R, G, N) post-RMSNorm, pre-bias      x, z: (b, l, H, P)
+      adt, dt: (b, H, l) fp32                            trap: (b, H, l) pre-sigmoid
+      C_bias, B_bias: (H, R, N)   mimo_*: (H, R, P)      angles: (b, l, H, A) raw
+      D: (H,)   norm_weight: (H*P,)
+      initial_states: (angle (b,H,A) f32, ssm (b,H,P,N) f32, k (b,R,H,N), v_raw (b,H,P))
+        exactly what the kernel returns with ``return_state=True`` -- ``k`` is the
+        last token's biased+rotated key and ``v_raw`` its *un-projected* x.
+
+    Per token t (upstream ``mamba3_MIMO_step_ref``, kernel rotate-half convention):
+      angle   += tanh(angles_t) * dt_t * pi
+      q, k     = rot(C_t + C_bias), rot(B_t + B_bias)          v_r = x_t * mimo_x[:, r]
+      alpha    = exp(adt_t);  gamma = sigmoid(trap_t) * dt_t;  beta = (1 - sigmoid) * dt_t * alpha
+      S        = alpha * S + beta * (k_{t-1}^T v_{t-1}) + gamma * (k_t^T v_t)   [sum over ranks]
+      y_r      = S q_r + D * v_r  ->  RMSNorm_P * w  * silu(z_t * mimo_z[:, r])
+      y        = sum_r mimo_o[:, r] * y_r
+    Returns ``(y (b, l, H, P) fp32, final states in the kernel's layout)``.
+    """
+    b, l, R, G, N = C.shape
+    H, P = x.shape[2], x.shape[3]
+    hpg = H // G
+    A = angles.shape[-1]
+    half = N // 2
+    if A > half:
+        raise ValueError(f"{A} rope angles exceed d_state/2 = {half}")
+
+    # head h reads group h // hpg  (kernel: i_h_qk = i_h // (H // G))
+    Cf = C.float().repeat_interleave(hpg, dim=3) + C_bias.float().permute(1, 0, 2)[None, None]
+    Bf = B.float().repeat_interleave(hpg, dim=3) + B_bias.float().permute(1, 0, 2)[None, None]
+    xf, zf = x.float(), z.float()
+    v_all = torch.einsum("blhp,hrp->blhrp", xf, mimo_x.float())
+    z_all = torch.einsum("blhp,hrp->blhrp", zf, mimo_z.float())
+    dt = dt.float()
+    sig = torch.sigmoid(trap.float())
+    alpha = torch.exp(adt.float())
+    gamma = sig * dt
+    beta = (1.0 - sig) * dt * alpha
+    w = norm_weight.float().view(H, P)
+    Df = D.float()
+    angles = angles.float()
+
+    if initial_states is not None:
+        angle_state, S, k_prev, v_prev_raw = initial_states
+        angle_state = angle_state.float().clone()
+        S = S.float().clone()
+        k_prev = k_prev.float().permute(0, 2, 1, 3)                       # (b, H, R, N)
+        v_prev = torch.einsum("bhp,hrp->bhrp", v_prev_raw.float(), mimo_x.float())
+    else:
+        angle_state = torch.zeros(b, H, A, dtype=torch.float32, device=x.device)
+        S = torch.zeros(b, H, P, N, dtype=torch.float32, device=x.device)
+        k_prev = torch.zeros(b, H, R, N, dtype=torch.float32, device=x.device)
+        v_prev = torch.zeros(b, H, R, P, dtype=torch.float32, device=x.device)
+
+    ys = []
+    for t in range(l):
+        q = Cf[:, t].permute(0, 2, 1, 3)                                  # (b, H, R, N)
+        k = Bf[:, t].permute(0, 2, 1, 3)
+        v = v_all[:, t]                                                   # (b, H, R, P)
+        angle_state = angle_state + torch.tanh(angles[:, t]) * dt[:, :, t, None] * math.pi
+        cos = angle_state.cos()[:, :, None, :]
+        sin = angle_state.sin()[:, :, None, :]
+        q = _rotate_half_partial(q, cos, sin)
+        k = _rotate_half_partial(k, cos, sin)
+        prev_kv = torch.einsum("bhrn,bhrp->bhpn", k_prev, v_prev)
+        curr_kv = torch.einsum("bhrn,bhrp->bhpn", k, v)
+        S = (
+            alpha[:, :, t, None, None] * S
+            + beta[:, :, t, None, None] * prev_kv
+            + gamma[:, :, t, None, None] * curr_kv
+        )
+        out = torch.einsum("bhpn,bhrn->bhrp", S, q) + Df[None, :, None, None] * v
+        out = out * torch.rsqrt(out.square().mean(-1, keepdim=True) + norm_eps)
+        out = out * w[None, :, None, :] * F.silu(z_all[:, t])
+        ys.append(torch.einsum("bhrp,hrp->bhp", out, mimo_o.float()))
+        k_prev, v_prev = k, v
+
+    y = torch.stack(ys, dim=1)
+    angle_state = torch.remainder(angle_state, 2 * math.pi)             # kernel's Final_Angle
+    finals = (angle_state, S, k_prev.permute(0, 2, 1, 3), xf[:, -1])
+    return y, finals
+
+
+class Mamba3GatedRMSNorm(nn.Module):
+    """Pre-gate grouped RMSNorm: ``norm_P(y) * w * silu(gate)`` per head of ``group_size``.
+
+    Same function as upstream's ``RMSNormGated(d_inner, group_size=headdim,
+    norm_before_gate=True)`` -- the "removed RMSNorm layer [that] is crucial for
+    long-context extrapolation" in the paper's hybrid ablation -- and the same
+    parameter name/shape (``norm.weight``, ``(d_inner,)``). The fused kernel applies
+    it inside the MIMO reduction (per rank, then ``mimo_o``), so this ``forward`` only
+    serves the non-fused rank-collapsed form; the reference path inlines the per-rank
+    version.
+    """
+
+    def __init__(self, hidden_size: int, group_size: int, eps: float = 1e-5):
+        super().__init__()
+        if hidden_size % group_size:
+            raise ValueError(f"{hidden_size} not divisible by group_size {group_size}")
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+        self.group_size = group_size
+
+    def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hs = hidden_states.float().unflatten(-1, (-1, self.group_size))
+        hs = hs * torch.rsqrt(hs.pow(2).mean(-1, keepdim=True) + self.eps)
+        hs = hs.flatten(-2) * self.weight.float() * F.silu(gate.float())
+        return hs.to(input_dtype)
+
+
+_KERNEL = None  # None: not tried yet; False: unavailable; else the callable
+
+
+def load_mimo_kernel(strict: bool = False):
+    """Lazily import upstream's fused ``mamba3_mimo`` (TileLang). Cached per process."""
+    global _KERNEL
+    if _KERNEL is None:
+        try:
+            from mamba_ssm.ops.tilelang.mamba3.mamba3_mimo import mamba3_mimo
+
+            _KERNEL = mamba3_mimo
+        except Exception as e:  # noqa: BLE001 - any import failure means "no kernel"
+            _KERNEL = False
+            msg = f"Mamba-3 MIMO kernel unavailable ({type(e).__name__}: {e}); using the torch reference"
+            if strict:
+                raise RuntimeError(msg) from e
+            warnings.warn(msg)
+    if strict and _KERNEL is False:
+        raise RuntimeError("Mamba-3 MIMO kernel unavailable (import failed earlier in this process)")
+    return _KERNEL
+
+
+class Qwen3Mamba3Mixer(nn.Module):
+    """Mamba-3 MIMO token mixer, upstream layout, inherited-attention init.
+
+    Set ``Qwen3Mamba3Mixer.force_reference = True`` to bypass the kernel (tests).
+    """
+
+    force_reference = False
+
+    def __init__(self, config: Qwen3Mamba3Config, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.mamba_n_heads
+        self.head_dim = config.mamba_d_head
+        self.ssm_state_size = config.mamba_d_state
+        self.n_groups = config.mamba_n_groups
+        self.mimo_rank = config.mamba_mimo_rank
+        self.intermediate_size = int(config.mamba_expand * config.hidden_size)
+        self.chunk_size = config.mamba_chunk_size
+        self.a_floor = config.mamba_a_floor
+        self.rotary_dim_divisor = mamba3_rotary_dim_divisor(config)
+        self.num_rope_angles = mamba3_num_rope_angles(config)
+        self.dt_min, self.dt_max = config.mamba_dt_min, config.mamba_dt_max
+        self.dt_init_floor = config.mamba_dt_init_floor
+        self.bc_bias_init = config.mamba_bc_bias_init
+        H, R, N, P, G = self.num_heads, self.mimo_rank, self.ssm_state_size, self.head_dim, self.n_groups
+        bc_eps = config.mamba_bc_norm_eps if config.mamba_bc_norm_eps is not None else config.rms_norm_eps
+
+        # [z | x | B | C | dd_dt | dd_A | trap | angles]  (upstream mamba3.py:112)
+        self.split_sizes = [
+            self.intermediate_size, self.intermediate_size,
+            N * G * R, N * G * R,
+            H, H, H,
+            self.num_rope_angles,
+        ]
+        self.in_proj = nn.Linear(self.hidden_size, sum(self.split_sizes), bias=False)
+        self.dt_bias = nn.Parameter(torch.ones(H))
+        self.dt_bias._no_weight_decay = True
+        self.B_bias = nn.Parameter(torch.zeros(H, R, N))
+        self.C_bias = nn.Parameter(torch.zeros(H, R, N))
+        # Qwen3RMSNorm rather than upstream's RMSNormGated: identical function (fp32
+        # rstd, weight multiply), and with eps = rms_norm_eps it IS the teacher's
+        # q_norm/k_norm, so inherited weights reproduce the teacher's normalized q/k
+        # exactly rather than to 1e-5.
+        self.B_norm = Qwen3RMSNorm(N, eps=bc_eps)
+        self.C_norm = Qwen3RMSNorm(N, eps=bc_eps)
+        self.mimo_x = nn.Parameter(torch.full((H, R, P), 1.0 / R))
+        self.mimo_z = nn.Parameter(torch.ones(H, R, P))
+        self.mimo_o = nn.Parameter(torch.full((H, R, P), 1.0 / R))
+        self.D = nn.Parameter(torch.ones(H))
+        self.D._no_weight_decay = True
+        self.norm = Mamba3GatedRMSNorm(self.intermediate_size, group_size=P, eps=config.mamba_norm_eps)
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    # ------------------------------------------------------------------ init
+    @torch.no_grad()
+    def init_ssm_parameters(self):
+        """Every non-Linear parameter, chosen so the mixer starts as its teacher head.
+
+        With the surgery's inheritance (C <- q_proj heads, B <- k_proj head, x <- v_proj
+        head, B/C_norm <- q/k_norm, out_proj <- o_proj) the block at init computes, per
+        head h with teacher kv group g and own rank r*:
+
+            y_h = norm( sum_s w_h(t, s) * <q_h(t), k_g(s)> * v_g(s) ) * silu(z)
+
+        i.e. the teacher's attention head with softmax replaced by a causal decay
+        ``w_h(t, s)`` -- a linearized copy, not a random SSM. Each line below is one
+        ingredient of that statement:
+
+        * ``dt_bias``: dtfix. dt ~ exp(U(log 1e-3, log 1e-1)) per head, floored at
+          1e-4, stored through the softplus inverse (upstream's own dt_min/dt_max and
+          the values Linearization_afei/DTFIX.md re-derived for Mamba2 after Bamba's
+          ``fill_(1.0)`` gave dt*|A| = 21 and a 4-token memory). Horizons of 10-1000
+          tokens across heads at A = -1.
+        * ``in_proj`` rows for ``dd_dt``, ``dd_A``, ``trap``, ``angles``: **zero**. Then
+          dt = softplus(dt_bias) is a per-head constant, A = -heavy_tail(0) = -1 for
+          every head, the trapezoid weight sigmoid(0) = 1/2, and the rotation angle is
+          0 -- the decay is a clean monotone distance kernel and no random RoPE
+          scrambles the inherited q.k geometry. All four keep a live gradient
+          (softplus', heavy_tail', sigmoid', sin' are nonzero at 0). The framework's
+          N(0, 0.02) rows would instead give every one of these an O(1) per-token
+          random swing at hidden 2560.
+        * ``B_bias``/``C_bias``: ``mamba_bc_bias_init`` (default 1.0, upstream's). It adds
+          ``sum(q) + sum(k) + N`` (N = 128) to every score at init, which looked like it
+          would swamp the inherited head; measured, the bias-1 arm beat bias 0 at every
+          stage-1 and 2a eval (see the config docstring), so the worry did not hold.
+        * ``mimo_x``: a random simplex over ranks per (h, p) -- softmax_r of N(0, 1) --
+          so ``sum_r v_r = x`` and the state ``sum_r k_r v_r^T = k x^T`` is *identical*
+          to upstream's uniform ``1/R`` at init (k is shared across ranks after
+          inheritance), while the per-rank values differ, which is what gives the
+          replicated B rows distinct gradients per rank instead of leaving the ranks
+          permanently tied.
+        * ``mimo_z``: ones (upstream): one gate for all ranks.
+        * ``mimo_o``: ``own_rank_matrix`` -- head h reads out the rank carrying teacher
+          q head h. Upstream's ``1/R`` would average the group's four heads into each.
+        * ``D``: ones (upstream). Under the pre-gate norm only the SSM:skip ratio
+          matters and the skip is small against O(|q.k|) readouts.
+        * ``norm.weight``: ones. Scale is re-fit by stage 1.
+        """
+        H, R, N, P, G = self.num_heads, self.mimo_rank, self.ssm_state_size, self.head_dim, self.n_groups
+        dev = self.dt_bias.device
+        dt = torch.exp(
+            torch.rand(H, dtype=torch.float32, device=dev)
+            * (math.log(self.dt_max) - math.log(self.dt_min))
+            + math.log(self.dt_min)
+        ).clamp(min=self.dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))  # softplus inverse
+        self.dt_bias.copy_(inv_dt.to(self.dt_bias.dtype))
+        self.B_bias.fill_(self.bc_bias_init)
+        self.C_bias.fill_(self.bc_bias_init)
+        self.mimo_x.copy_(
+            torch.softmax(torch.randn(H, R, P, dtype=torch.float32, device=dev), dim=1).to(self.mimo_x.dtype)
+        )
+        self.mimo_z.fill_(1.0)
+        self.mimo_o.copy_(own_rank_matrix(H, G, R).to(dev)[:, :, None].expand(H, R, P).to(self.mimo_o.dtype))
+        self.D.fill_(1.0)
+        self.norm.weight.fill_(1.0)
+        self.B_norm.weight.fill_(1.0)
+        self.C_norm.weight.fill_(1.0)
+        scalar_rows = 2 * self.intermediate_size + 2 * N * G * R
+        self.in_proj.weight[scalar_rows:].zero_()
+
+    # --------------------------------------------------------------- forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params=None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        b, l, _ = hidden_states.shape
+        H, R, N, P, G = self.num_heads, self.mimo_rank, self.ssm_state_size, self.head_dim, self.n_groups
+        dtype = hidden_states.dtype
+
+        # 2D padding mask (Bamba's apply_mask_to_padding_states): a zeroed input gives
+        # z = x = B = C = 0 and dd_* = 0, so a pad token writes nothing, rotates
+        # nothing, and only decays the state by exp(A * softplus(dt_bias)).
+        if attention_mask is not None and attention_mask.dim() == 2 and attention_mask.shape[1] == l:
+            hidden_states = hidden_states * attention_mask[:, :, None].to(dtype)
+
+        proj = self.in_proj(hidden_states)
+        z, x, B, C, dd_dt, dd_A, trap, angles = torch.split(proj, self.split_sizes, dim=-1)
+        z = z.view(b, l, H, P)
+        x = x.view(b, l, H, P)
+        B = self.B_norm(B.view(b, l, R, G, N))          # rank-major "(r g n)", as upstream
+        C = self.C_norm(C.view(b, l, R, G, N))
+
+        dt = F.softplus(dd_dt.float() + self.dt_bias.float())                      # (b, l, H)
+        A = torch.clamp(-heavy_tail_activation(dd_A.float()), max=-self.a_floor)   # (b, l, H)
+        dt_t = dt.transpose(1, 2)                                                  # (b, H, l)
+        adt_t = (A * dt).transpose(1, 2)
+        trap_t = trap.transpose(1, 2)                                              # raw, pre-sigmoid
+        angles = angles.float().unsqueeze(2).expand(b, l, H, self.num_rope_angles)  # raw, pre-tanh
+
+        states = None
+        warm = False
+        if cache_params is not None:
+            i = self.layer_idx
+            states = (
+                cache_params.angle_states[i], cache_params.ssm_states[i],
+                cache_params.k_states[i], cache_params.v_states[i],
+            )
+            warm = bool(cache_params.has_previous_state)
+
+        kernel = None
+        if x.is_cuda and not warm and not self.force_reference:
+            kernel = load_mimo_kernel() or None
+
+        if kernel is not None:
+            out = kernel(
+                Q=C, K=B, V=x,
+                ADT=adt_t, DT=dt_t, Trap=trap_t.to(dtype),
+                Q_bias=self.C_bias.float(), K_bias=self.B_bias.float(),
+                MIMO_V=self.mimo_x.float(), MIMO_Z=self.mimo_z.float(), MIMO_Out=self.mimo_o.float(),
+                Angles=angles, D=self.D.float(), Z=z,
+                chunk_size=self.chunk_size,
+                rotary_dim_divisor=self.rotary_dim_divisor,
+                dtype=dtype,
+                return_state=states is not None,
+                fuse_pregate_headwise_rms_norm=True,
+                outproj_norm_weight=self.norm.weight.float(),
+                outproj_norm_eps=self.norm.eps,
+            )
+            if states is not None:
+                y, angle_f, ssm_f, k_f, v_f = out
+                states[0].copy_(angle_f)
+                states[1].copy_(ssm_f)
+                states[2].copy_(k_f)
+                states[3].copy_(v_f)
+            else:
+                y = out
+        else:
+            y, finals = mamba3_mimo_reference(
+                C, B, x, z, adt_t, dt_t, trap_t,
+                self.C_bias, self.B_bias, self.mimo_x, self.mimo_z, self.mimo_o,
+                angles, self.D, self.norm.weight, self.norm.eps,
+                initial_states=states if warm else None,
+            )
+            if states is not None:
+                for dst, src in zip(states, finals):
+                    dst.copy_(src.to(dst.dtype))
+
+        y = y.reshape(b, l, self.intermediate_size)
+        return self.out_proj(y.to(dtype))
+
+
+class Qwen3Mamba3HybridCache(HybridMambaAttentionDynamicCache):
+    """Bamba's hybrid cache with Mamba-3's four recurrent states per converted layer.
+
+    Per converted layer (kernel ``return_state`` layout, upstream ``allocate_inference_cache``):
+      ``angle_states``: (B, H, A) fp32 -- cumulative rotation angle, mod 2 pi
+      ``ssm_states``:   (B, H, P, N) fp32
+      ``k_states``:     (B, R, H, N) model dtype -- last biased+rotated key
+      ``v_states``:     (B, H, P) model dtype -- last raw x (pre-``mimo_x``)
+    ``conv_states`` are kept as ``(B, 0)`` placeholders (no conv in Mamba-3) so the
+    parent's ``reorder_cache`` keeps working; the three fixes from the Mamba2 cache
+    (int widths, ``get_seq_length`` on an unwritten cache, ``get_mask_sizes``) carry
+    over.
+    """
+
+    def __init__(self, config: Qwen3Mamba3Config, batch_size: int, dtype=torch.float16, device=None):
+        self.layers_block_type = config.layers_block_type
+        self.has_previous_state = False
+        H, P, N = config.mamba_n_heads, config.mamba_d_head, config.mamba_d_state
+        R, A = config.mamba_mimo_rank, mamba3_num_rope_angles(config)
+
+        def placeholder():
+            return torch.tensor([[]] * batch_size, device=device)
+
+        self.conv_states, self.ssm_states, self.transformer_layers = [], [], []
+        self.angle_states, self.k_states, self.v_states = [], [], []
+        for i in range(config.num_hidden_layers):
+            self.conv_states.append(placeholder())
+            if self.layers_block_type[i] == "mamba":
+                self.angle_states.append(torch.zeros(batch_size, H, A, device=device, dtype=torch.float32))
+                self.ssm_states.append(torch.zeros(batch_size, H, P, N, device=device, dtype=torch.float32))
+                self.k_states.append(torch.zeros(batch_size, R, H, N, device=device, dtype=dtype))
+                self.v_states.append(torch.zeros(batch_size, H, P, device=device, dtype=dtype))
+            else:
+                self.angle_states.append(placeholder())
+                self.ssm_states.append(placeholder())
+                self.k_states.append(placeholder())
+                self.v_states.append(placeholder())
+                self.transformer_layers.append(i)
+        self.key_cache = [placeholder() for _ in range(config.num_hidden_layers)]
+        self.value_cache = [placeholder() for _ in range(config.num_hidden_layers)]
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        super().reorder_cache(beam_idx)
+        for states in (self.angle_states, self.k_states, self.v_states):
+            for i in range(len(states)):
+                states[i] = states[i].index_select(0, beam_idx.to(states[i].device))
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        if layer_idx not in self.transformer_layers:
+            layer_idx = self.transformer_layers[0]
+        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx].numel() == 0:
+            return 0
+        return self.key_cache[layer_idx].shape[-2]
+
+    def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:
+        query_length = cache_position.shape[0]
+        return self.get_seq_length(layer_idx) + query_length, 0
+
+
+class Qwen3Mamba3LinearDecoderLayer(GradientCheckpointingLayer):
+    """Pre-norm decoder layer with the Mamba-3 MIMO mixer under ``linear_attn``."""
+
+    def __init__(self, config: Qwen3Mamba3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layer_type = self.attention_type = "linear_attention"
+        self.linear_attn = Qwen3Mamba3Mixer(config, layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings=None,  # unused: the mixer's rotation is data-dependent
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.linear_attn(
+            hidden_states=hidden_states,
+            cache_params=past_key_values,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
+
+class Qwen3Mamba3PreTrainedModel(PreTrainedModel):
+    config_class = Qwen3Mamba3Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer", "Qwen3Mamba3LinearDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _is_stateful = True
+
+    def _init_weights(self, module):
+        # Generic init first (Linear N(0, 0.02), *RMSNorm* weights -> 1), then the
+        # mixer's own parameters. ``nn.Module.apply`` visits children before the
+        # parent, so ``in_proj`` is already drawn when the mixer zeroes its scalar rows.
+        super()._init_weights(module)
+        if isinstance(module, Qwen3Mamba3Mixer):
+            module.init_ssm_parameters()
+
+
+class Qwen3Mamba3Model(Qwen3Mamba3PreTrainedModel):
+    def __init__(self, config: Qwen3Mamba3Config):
+        super().__init__(config)
+        assert config.layer_types is not None and len(config.layer_types) == config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DecoderLayer(config, i)
+                if config.layer_types[i] == "full_attention"
+                else Qwen3Mamba3LinearDecoderLayer(config, i)
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if self.gradient_checkpointing and self.training:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = Qwen3Mamba3HybridCache(
+                config=self.config,
+                batch_size=inputs_embeds.shape[0],
+                dtype=inputs_embeds.dtype,
+                device=inputs_embeds.device,
+            )
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+        linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers:
+            layer_mask = (
+                linear_attn_mask if decoder_layer.attention_type == "linear_attention" else causal_mask
+            )
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=layer_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values if use_cache else None,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        # The mixer takes the recurrent (warm) path only once this flag is up; the
+        # model raises it after the first cached forward (same contract as Bamba).
+        if use_cache and past_key_values is not None and not past_key_values.has_previous_state:
+            past_key_values.has_previous_state = True
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+    def _update_linear_attn_mask(self, attention_mask, cache_position):
+        linear_attn_mask = attention_mask
+        if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)):
+            linear_attn_mask = None
+        return linear_attn_mask
+
+
+class Qwen3Mamba3ForCausalLM(Qwen3Mamba3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: Qwen3Mamba3Config):
+        super().__init__(config)
+        self.model = Qwen3Mamba3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep=0,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )
+
+
+def register():
+    """Make qwen3_mamba3 checkpoints loadable through the Auto* classes."""
+    from transformers import AutoConfig, AutoModelForCausalLM
+
+    try:
+        AutoConfig.register("qwen3_mamba3", Qwen3Mamba3Config)
+    except ValueError:
+        pass  # already registered in this process
+    AutoModelForCausalLM.register(Qwen3Mamba3Config, Qwen3Mamba3ForCausalLM, exist_ok=True)
+
+
+try:  # convenience for direct imports; auto_map already covers the
+    register()  # remote-code path, so a registry collision is not fatal here
+except Exception:  # pragma: no cover
+    pass