ssm.py

"""Fallback padding-safe recurrent SSM scan.

This is not a fast Mamba kernel. It is a deterministic reference path for
correctness tests, mask semantics, and future kernel equivalence checks.
"""

from __future__ import annotations

import torch
from torch import nn


class MaskedScanSSM(nn.Module):
    def __init__(self, hidden_size: int, state_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        self.state_size = state_size
        self.in_proj = nn.Linear(hidden_size, state_size)
        self.state_proj = nn.Linear(state_size, state_size, bias=False)
        self.out_proj = nn.Linear(state_size, hidden_size)
        self.gate = nn.Linear(hidden_size, hidden_size)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        segment_ids: torch.Tensor | None = None,
        reset_on_pad: bool = True,
        reset_on_segment: bool = True,
    ) -> torch.Tensor:
        batch, length, _ = hidden_states.shape
        state = hidden_states.new_zeros(batch, self.state_size)
        outputs = []
        projected = self.in_proj(hidden_states)
        for idx in range(length):
            valid = attention_mask[:, idx].unsqueeze(-1)
            if reset_on_segment and segment_ids is not None and idx > 0:
                same_segment = (segment_ids[:, idx] == segment_ids[:, idx - 1]).unsqueeze(-1)
                state = torch.where(same_segment, state, torch.zeros_like(state))
            if reset_on_pad:
                state = torch.where(valid, state, torch.zeros_like(state))
            proposed = torch.tanh(projected[:, idx] + self.state_proj(state))
            state = torch.where(valid, proposed, torch.zeros_like(state) if reset_on_pad else state)
            out = self.out_proj(state) * torch.sigmoid(self.gate(hidden_states[:, idx]))
            outputs.append(torch.where(valid, out, torch.zeros_like(out)))
        return torch.stack(outputs, dim=1)