feat: Early refactor merge (#38)

* feat: Add chunk param selection heuristic for xl_chunk size and align last states format (#33)

* chore: rename.

* wip: change state shape.

* feat: add heuristic kernel parameter selection to xl_chunk_size kernel.

* chore: update doc.

* chore: improve readability.

* feat: Add all sequence length kernel dispatcher (#34)

* chore: rename.

* wip: change state shape.

* feat: add heuristic kernel parameter selection to xl_chunk_size kernel.

* wip: add arbitrary sequence length wrapper + test. Test fails atm.

* fix: arbitrary sequence length wrapper.

* feat: make state dtype of trition step fused kernel configurable.

* feat: add test for limit + xl chunk size kernels.

* chore: improve doc.

* fix: native sequence return states.

* feat: enable single step in arbitrary sequence length.

* chore: add mlstm sequence kernel registry.

* chore: fix typo.

* Merge branch 'refactor_further' into add_all_sequence_length_kernel_dispatcher

* chore: prettify

* chore: use tl.exp instead of tl.exp2 * const

* chore: improve error message.

* chore: control logging only via logger.

* chore: fix merge

* chore: simplify kernel call logic.

* chore: add doc strings.

* feat: Add pytorch backend module (#35)

* chore: rename.

* wip: change state shape.

* feat: add heuristic kernel parameter selection to xl_chunk_size kernel.

* wip: add arbitrary sequence length wrapper + test. Test fails atm.

* fix: arbitrary sequence length wrapper.

* feat: make state dtype of trition step fused kernel configurable.

* feat: add test for limit + xl chunk size kernels.

* chore: improve doc.

* fix: native sequence return states.

* feat: enable single step in arbitrary sequence length.

* chore: add mlstm sequence kernel registry.

* chore: fix typo.

* feat: add preliminary backend module.

* chore: support different state dtype in triton step kernel. Store outputs only from a single thread block, i.e. avoid duplicate store.

* feat: fully support state dtype in recurrent sequence loop.

* fix: handle the case where target chunk size is smaller than default block size.

* feat: add pad zeros kernel wrapper.

* feat: integrate pad with zeros into backend module.

* chore: add doc to mlstm backend config.

* chore: add doc to forward pass.

* fix: typo. add doc.

* chore: improve error message.

* chore: clean up backend module.

* fix: merge

* chore: change number format.

---------

Co-authored-by: Maximilian Beck <[email protected]>

* chore: remove deprecated argument.

---------

Co-authored-by: Maximilian Beck <[email protected]>

Author: maximilianmbeck

.vscode/launch.json

@@ -17,7 +17,7 @@
             "justMyCode": false,
             "args": [
                 // "-cn",
-                "${workspaceFolder}/tests/test_mlstm/test_parallel/test_parallel_torch.py",
+                "${workspaceFolder}/tests/torch/test_arbitrary_sequence_length.py",
             ],
             "env": {
                 "CUDA_VISIBLE_DEVICES": "0",

.vscode/settings.json

@@ -7,7 +7,7 @@
       "source.organizeImports": "explicit"
     }
   },
-  "editor.formatOnSave": true,
+  "editor.formatOnSave": false,
   "isort.args": ["--profile", "ruff"],
   "files.watcherExclude": {
     "outputs/**": true,

README.md

@@ -8,19 +8,19 @@ In this repository we collect clean implementations of the different mLSTM formu
 def mlstm_interface(
     q: torch.Tensor, # (B, NH, S, DHQK)
     k: torch.Tensor, # (B, NH, S, DHQK)
-    v: torch.Tensor, # (B, NH, S, DHV)
+    v: torch.Tensor, # (B, NH, S, DHHV)
     i: torch.Tensor, # (B, NH, S)
     f: torch.Tensor, # (B, NH, S)
-    c_initial: torch.Tensor = None, # (B, NH, DHQK, DHV)
+    c_initial: torch.Tensor = None, # (B, NH, DHQK, DHHV)
     n_initial: torch.Tensor = None, # (B, NH, DHQK)
-    m_initial: torch.Tensor = None, # (B, NH) # TODO change the shape of this to (B, NH, 1)
+    m_initial: torch.Tensor = None, # (B, NH, 1)
     return_last_states: bool = False,
     eps: float = 1e-6,
-    autocast_kernel_dtype: torch.dtype = torch.float16,
+    autocast_kernel_dtype: torch.dtype = torch.bfloat16,
     chunk_size: int = 64,
     **kwargs,
 ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
-    # (B, NH, S, DHV) | ((B, NH, S, DHV), ((B, NH, DHQK, DHV), (B, NH, DHQK), (B, NH)))
+    # (B, NH, S, DHHV) | ((B, NH, S, DHHV), ((B, NH, DHQK, DHHV), (B, NH, DHQK), (B, NH)))
     """
     Returns:
         torch.Tensor: matH outputs (no n and m values, no last states)
@@ -35,17 +35,17 @@ def mlstm_interface(
 def mlstm_step_interface(
     q: torch.Tensor,  # (B, NH, DHQK)
     k: torch.Tensor,  # (B, NH, DHQK)
-    v: torch.Tensor,  # (B, NH, DHV)
+    v: torch.Tensor,  # (B, NH, DHHV)
     i: torch.Tensor,  # (B, NH, 1)
     f: torch.Tensor,  # (B, NH, 1)
-    c: torch.Tensor,  # (B, NH, DHQK, DHV)
+    c: torch.Tensor,  # (B, NH, DHQK, DHHV)
     n: torch.Tensor,  # (B, NH, DHQK)
     m: torch.Tensor,  # (B, NH, 1)
     eps: float = 1e-6,
     **kwargs,
 ) -> tuple[
     torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-]:  # vecH, (matC_state_new (B, NH, DHQK, DHV), vecN_state_new (B, NH, DHQK), vecM_state_new (B, NH, 1))
+]:  # vecH, (matC_state_new (B, NH, DHQK, DHHV), vecN_state_new (B, NH, DHQK), vecM_state_new (B, NH, 1))
 ```
 
 ## Kernel variants

mlstm_kernels/torch/__init__.py

@@ -11,9 +11,13 @@ def _create_module_sequence_backend_registry() -> dict[str, dict[str, Callable]]
     }
     return module_backend_registry
 
+
 def get_available_mlstm_kernels() -> list[str]:
     """
-    Get a list of available mlstm sequence kernel names.
+    Get a list of available mlstm sequence kernels.
+    These kernels process a sequence in the parallel or chunkwise parallel mode of the mLSTM. 
+    They do not support arbitrary sequence lengths. 
+    They are used for training and prefill processing during inference of the mLSTM during.
     """
     module_backend_registry = _create_module_sequence_backend_registry()
 
@@ -24,11 +28,6 @@ def get_available_mlstm_kernels() -> list[str]:
     ]
     return backend_names
 
-def get_available_mlstm_step_kernels() -> list[str]:
-    from .recurrent import registry_step as mlstm_recurrent_step_registry
-    backend_names = list(mlstm_recurrent_step_registry.keys())
-    return backend_names
-
 
 def get_mlstm_kernel(name: str) -> Callable:
     """
@@ -54,12 +53,22 @@ def get_mlstm_kernel(name: str) -> Callable:
 
     if backend_name not in module_backend_registry[module_name]:
         raise ValueError(
-            f"Unknown backend name: {backend_name}. Available backend names: {list(module_backend_registry[module_name].keys())}"
+            f"Unknown mlstm kernel backend name: {backend_name}. Available backend names: {list(module_backend_registry[module_name].keys())}"
         )
 
     return module_backend_registry[module_name][backend_name]
 
 
+def get_available_mlstm_step_kernels() -> list[str]:
+    """Returns the available mlstm step kernels. 
+    These kernels can be used to compute a single time step of the mLSTM, i.e. for generation.
+    """
+    from .recurrent import registry_step as mlstm_recurrent_step_registry
+
+    backend_names = list(mlstm_recurrent_step_registry.keys())
+    return backend_names
+
+
 def get_mlstm_step_kernel(name: str) -> Callable:
     """
     Get a mlstm step kernel function by name.
@@ -73,7 +82,36 @@ def get_mlstm_step_kernel(name: str) -> Callable:
 
     if name not in mlstm_recurrent_step_registry:
         raise ValueError(
-            f"Unknown backend name: {name}. Available backend names: {list(mlstm_recurrent_step_registry.keys())}"
+            f"Unknown step kernel backend name: {name}. Available backend names: {list(mlstm_recurrent_step_registry.keys())}"
+        )
+
+    return mlstm_recurrent_step_registry[name]
+
+
+def get_available_mlstm_sequence_kernels() -> list[str]:
+    """Returns the available mlstm sequence kernels.
+    These kernels process a sequence in the recurrent mode of the mLSTM and hence support any sequence length.
+    """
+    from .recurrent import registry_sequence as mlstm_recurrent_sequence_registry
+
+    backend_names = list(mlstm_recurrent_sequence_registry.keys())
+    return backend_names
+
+
+def get_mlstm_sequence_kernel(name: str) -> Callable:
+    """
+    Get a mlstm sequence kernel function by name.
+
+    Naming convention:
+    name = "<backend_name>"
+
+    backend_name: The name of the kernel function as defined in the registry in the __init__.py file of the module.
+    """
+    from .recurrent import registry_sequence as mlstm_recurrent_sequence_registry
+
+    if name not in mlstm_recurrent_sequence_registry:
+        raise ValueError(
+            f"Unknown backend name: {name}. Available backend names: {list(mlstm_recurrent_sequence_registry.keys())}"
         )
 
-    return mlstm_recurrent_step_registry[name]
+    return mlstm_recurrent_sequence_registry[name]

mlstm_kernels/torch/backend_module.py

@@ -0,0 +1,205 @@
+from dataclasses import dataclass
+from functools import partial
+from typing import Literal
+
+import torch
+from torch import nn
+
+from . import (
+    get_mlstm_kernel,
+    get_mlstm_sequence_kernel,
+    get_mlstm_step_kernel,
+)
+from .kernel_wrappers import (
+    wrap_chunkwise__arbitrary_sequence_length,
+    wrap_chunkwise__pad_zeros,
+)
+
+ChunkwiseKernelType = Literal[
+    "chunkwise--native_autograd",
+    "chunkwise--native_custbw",
+    "chunkwise--triton_limit_chunk",
+    "chunkwise--triton_xl_chunk",
+    "parallel--native_autograd",
+    "parallel--native_custbw",
+    "parallel--native_stablef_autograd",
+    "parallel--native_stablef_custbw",
+    "parallel--triton_limit_headdim",
+]
+SequenceKernelType = Literal[
+    "native_sequence__native", "native_sequence__triton_step_fused"
+]
+StepKernelType = Literal["native", "triton_fused"]
+
+DtypeType = Literal["float32", "bfloat16", "float16"]
+
+BackendModeType = Literal["train", "train_with_padding", "inference"]
+
+
+@dataclass
+class mLSTMBackendConfig:
+    chunkwise_kernel: ChunkwiseKernelType = "chunkwise--native_autograd"
+    """The chunkwise kernel to use for chunkwise parallel processing of the sequence.
+    This kernel is used for training.
+    Also supports fully parallel (i.e. quadratic) backends for comparison.
+    """
+    sequence_kernel: SequenceKernelType = "native_sequence__native"
+    """The sequence kernel to use for processing sequneces step-by-step.
+    Used only for parts of the prefill sequence in inference mode.
+    """
+    step_kernel: StepKernelType = "native"
+    """The step kernel to use for processing a single step.
+    Used for generation in inference mode.
+    """
+    mode: BackendModeType = "train"
+    """The mode of operation for the backend. Determines how the `forward` method behaves.
+    """
+    chunk_size: int = 64
+    """The chunk size of the chunkwise kernel.
+    If the mode is 'train_with_padding', this is the inputs are padded to multiples of this size.
+    """
+    return_last_states: bool = True
+    """Whether to return the last states of the sequence in training mode.
+    Inference mode always returns the last states.
+    """
+    autocast_kernel_dtype: DtypeType = "bfloat16"
+    """The dtype to use for autocast behavior in the kernel.
+    If autocast is enabled all inputs are cast to this dtype before the kernel is called.
+    """
+    eps: float = 1e-6
+    """Epsilon value for numerical stability in the kernel."""
+    inference_state_dtype: DtypeType = "float32"
+    """The dtype to use for the state tensors in inference mode."""
+
+    def __post_init__(self):
+        if self.return_last_states and "parallel" in self.chunkwise_kernel:
+            raise ValueError(
+                "return_last_states=True is not supported with parallel kernels."
+            )
+        if self.return_last_states and self.mode == "train_with_padding":
+            raise ValueError(
+                "return_last_states=True is not supported with train_with_padding mode."
+            )
+
+
+class mLSTMBackend(nn.Module):
+    """mLSTM Backend Module for PyTorch.
+
+    This module wraps the mLSTM kernels and provides a high-level interface for training and inference.
+    """
+
+    config_class = mLSTMBackendConfig
+
+    def __init__(self, config: mLSTMBackendConfig):
+        super().__init__()
+        self.config = config
+        self.chunkwise_kernel_fn = get_mlstm_kernel(config.chunkwise_kernel)
+        self.sequence_kernel_fn = get_mlstm_sequence_kernel(config.sequence_kernel)
+        self.step_kernel_fn = get_mlstm_step_kernel(config.step_kernel)
+
+        self._inference_fn = partial(
+            wrap_chunkwise__arbitrary_sequence_length,
+            mlstm_chunkwise_kernel=self.chunkwise_kernel_fn,
+            mlstm_sequence_kernel=partial(
+                self.sequence_kernel_fn,
+                dtype_state=getattr(torch, config.inference_state_dtype),
+            ),
+            mlstm_step_kernel=partial(
+                self.step_kernel_fn,
+                dtype_state=getattr(torch, config.inference_state_dtype),
+            ),
+            chunk_size=config.chunk_size,
+            eps=config.eps,
+            autocast_kernel_dtype=getattr(torch, config.autocast_kernel_dtype),
+            return_last_states=True,
+        )
+
+        train_kernel_fn = partial(
+            self.chunkwise_kernel_fn,
+            autocast_kernel_dtype=getattr(torch, config.autocast_kernel_dtype),
+            eps=config.eps,
+            chunk_size=config.chunk_size,
+        )
+        if "with_padding" in config.mode:
+            train_kernel_fn = partial(
+                wrap_chunkwise__pad_zeros, mlstm_chunkwise_kernel=train_kernel_fn
+            )
+        self._train_fn = train_kernel_fn
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        i: torch.Tensor,
+        f: torch.Tensor,
+        c_initial: torch.Tensor = None,
+        n_initial: torch.Tensor = None,
+        m_initial: torch.Tensor = None,
+        return_last_states: bool = None,
+        mode: Literal["train", "inference"] = None,
+    ) -> (
+        torch.Tensor
+        | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
+    ):
+        """Forward pass of the mLSTM backend.
+
+        Depending on the configured mode, this method will call the appropriate kernel function.
+
+        Args:
+            q: The query tensor of shape (B, NH, S, DHQK).
+            k: The key tensor of shape (B, NH, S, DHQK).
+            v: The value tensor of shape (B, NH, S, DHHV).
+            i: The input gate preactivation tensor of shape (B, NH, S).
+            f: The forget gate preactivation tensor of shape (B, NH, S).
+            c_initial: The initial cell state tensor of shape (B, NH, DHQK, DHHV).
+                                                Defaults to None.
+            n_initial: The initial hidden state tensor of shape (B, NH, DHQK). Defaults to None.
+            m_initial: The initial memory tensor of shape (B, NH, 1). Defaults to None.
+            return_last_states: Whether to return the last states of the sequence. Defaults to None.
+                                                If None, the value from the config is used.
+
+        Returns:
+            hidden states of shape (B, NH, S, DHHV)
+            hidden states and last states the last states are the cell state c (B, NH, DHQK, DHHV),
+            the normalizer state n (B, NH, DHQK), and the max state m (B, NH, 1)
+        """
+        if mode is None:
+            mode = self.config.mode
+
+        if "train" in mode:
+            if return_last_states is None:
+                return_last_states = self.config.return_last_states
+
+            if self.config.mode == "train_with_padding":
+                assert not return_last_states, "return_last_states=True is not supported with train_with_padding mode."
+
+            return self._train_fn(
+                q=q,
+                k=k,
+                v=v,
+                i=i,
+                f=f,
+                c_initial=c_initial,
+                n_initial=n_initial,
+                m_initial=m_initial,
+                return_last_states=return_last_states,
+            )
+
+        elif "inference" in mode:
+            # inference mode always returns the last states
+            return self._inference_fn(
+                q=q,
+                k=k,
+                v=v,
+                i=i,
+                f=f,
+                c_initial=c_initial,
+                n_initial=n_initial,
+                m_initial=m_initial,
+            )
+        else:
+            raise ValueError(f"Unknown mode: {self.config.mode}")
+
+    def extra_repr(self) -> str:
+        return f"{self.config}"