Fix Windows inference: SDPA math fallback, FLA lazy-import patches, clearer install docs

README.md

@@ -144,12 +144,35 @@ Monostich-2 needs **GatedDeltaNet-2** from a pinned `flash-linear-attention` com
 
 ```bash
 pip install safetensors tokenizers huggingface_hub einops numpy
+```
+
+#### Linux / WSL (recommended)
 
-# Pin used for Monostich-2 training / inference
+```bash
 pip install --no-deps "git+https://github.com/fla-org/flash-linear-attention.git@cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
 ```
 
-> If the git install fails on Windows, install [Git for Windows](https://git-scm.com/download/win) first, or clone the repo and `pip install --no-deps -e flash-linear-attention` after checking out that commit.
+#### Windows (native)
+
+Stock FLA + Triton 3.7 (bundled with recent PyTorch wheels) crashes on import.
+Use UTF-8 mode, install editable from a clone, then apply the Hub patches:
+
+```powershell
+$env:PYTHONUTF8 = "1"
+pip install safetensors tokenizers huggingface_hub einops numpy
+
+git clone https://github.com/fla-org/flash-linear-attention.git
+cd flash-linear-attention
+git checkout cbb0a72efb55c18ca0ef4f298298317573ad2cb3
+
+hf download kerzgrr/Monostich-2-base --include "windows_fla_patches/**" --local-dir "$env:TEMP\m2patches"
+Copy-Item -Recurse -Force "$env:TEMP\m2patches\windows_fla_patches\fla\*" .\fla\
+
+pip install --no-build-isolation --no-deps -e .
+cd ..
+```
+
+> Alternative that Just Works: run inference inside **WSL** with the Linux install path above.
 
 ### 5) Download the inference script
 

inference.py

@@ -239,18 +239,30 @@ def main() -> int:
         sys.path.insert(0, str(package_root))
 
     try:
-        from tiny_gdn import TinyGDNConfig, TinyGDNForCausalLM
-    except ImportError as error:
+        from fla.layers.gdn2 import GatedDeltaNet2  # noqa: F401
+    except Exception as error:  # noqa: BLE001
         print(
-            "Could not import tiny_gdn / GatedDeltaNet-2.\n"
+            "Could not import GatedDeltaNet-2 from flash-linear-attention.\n"
             "Install dependencies first (see README):\n"
             "  pip install torch safetensors tokenizers huggingface_hub einops\n"
+            f"  # Linux/WSL:\n"
             f"  pip install --no-deps git+https://github.com/fla-org/flash-linear-attention.git@{FLA_COMMIT}\n"
+            "  # Windows: clone that commit, apply windows_fla_patches/ from this repo, then\n"
+            "  #   pip install --no-build-isolation --no-deps -e .\n"
             f"Import error: {error}",
             file=sys.stderr,
         )
         return 1
 
+    try:
+        from tiny_gdn import TinyGDNConfig, TinyGDNForCausalLM
+    except ImportError as error:
+        print(
+            f"Could not import tiny_gdn: {error}",
+            file=sys.stderr,
+        )
+        return 1
+
     weights_path = _download("model.safetensors", local_dir)
     tok_path = _download("tokenizer.json", local_dir)
     cfg_path = _download("config.json", local_dir)

tiny_gdn/model.py

@@ -15,7 +15,9 @@ from torch.utils.checkpoint import checkpoint
 from tiny_gdn.config import TinyGDNConfig
 
 try:
-    from fla.layers import GatedDeltaNet2
+    # Import the module directly — `from fla.layers import GatedDeltaNet2`
+    # executes layers/__init__.py and eagerly loads every attention kernel.
+    from fla.layers.gdn2 import GatedDeltaNet2
 except ImportError as import_error:
     GatedDeltaNet2 = None
     FLA_IMPORT_ERROR: ImportError | None = import_error
@@ -198,15 +200,19 @@ class GatedGroupedQueryAttention(nn.Module):
             "is_causal": sdpa_mask is None,
             "enable_gqa": True,
         }
-        if query.is_cuda:
-            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
-                attention_output = F.scaled_dot_product_attention(
-                    query,
-                    key,
-                    value,
-                    **sdpa_options,
-                )
-        else:
+        # Prefer Flash / mem-efficient when available; fall back to MATH for
+        # Windows PyTorch builds that ship without FlashAttention kernels.
+        backends = (
+            [
+                SDPBackend.FLASH_ATTENTION,
+                SDPBackend.EFFICIENT_ATTENTION,
+                SDPBackend.CUDNN_ATTENTION,
+                SDPBackend.MATH,
+            ]
+            if query.is_cuda
+            else [SDPBackend.MATH]
+        )
+        with sdpa_kernel(backends):
             attention_output = F.scaled_dot_product_attention(
                 query,
                 key,

windows_fla_patches/fla/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Keep package import light. Eagerly importing fla.layers pulls every Triton
+# kernel and breaks on Windows + Triton 3.7.
+
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+__version__ = "0.5.2"
+__all__: list[str] = []

windows_fla_patches/fla/layers/__init__.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy layer exports — avoid compiling every Triton kernel at import time.
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, tuple[str, str]] = {
+    "ABCAttention": (".abc", "ABCAttention"),
+    "Attention": (".attn", "Attention"),
+    "BasedLinearAttention": (".based", "BasedLinearAttention"),
+    "BitAttention": (".bitattn", "BitAttention"),
+    "Comba": (".comba", "Comba"),
+    "DeltaNet": (".delta_net", "DeltaNet"),
+    "DeltaFormerAttention": (".deltaformer", "DeltaFormerAttention"),
+    "ForgettingAttention": (".forgetting_attn", "ForgettingAttention"),
+    "GatedDeltaNet": (".gated_deltanet", "GatedDeltaNet"),
+    "GatedDeltaProduct": (".gated_deltaproduct", "GatedDeltaProduct"),
+    "GatedDeltaNet2": (".gdn2", "GatedDeltaNet2"),
+    "GatedLinearAttention": (".gla", "GatedLinearAttention"),
+    "GatedSlotAttention": (".gsa", "GatedSlotAttention"),
+    "HGRNAttention": (".hgrn", "HGRNAttention"),
+    "HGRN2Attention": (".hgrn2", "HGRN2Attention"),
+    "KimiDeltaAttention": (".kda", "KimiDeltaAttention"),
+    "LightNetAttention": (".lightnet", "LightNetAttention"),
+    "LinearAttention": (".linear_attn", "LinearAttention"),
+    "LogLinearMamba2": (".log_linear_mamba2", "LogLinearMamba2"),
+    "Mamba": (".mamba", "Mamba"),
+    "Mamba2": (".mamba2", "Mamba2"),
+    "Mamba3": (".mamba3", "Mamba3"),
+    "MesaNet": (".mesa_net", "MesaNet"),
+    "MultiheadLatentAttention": (".mla", "MultiheadLatentAttention"),
+    "MoBA": (".moba", "MoBA"),
+    "MomAttention": (".mom", "MomAttention"),
+    "MultiScaleRetention": (".multiscale_retention", "MultiScaleRetention"),
+    "NativeSparseAttention": (".nsa", "NativeSparseAttention"),
+    "Parallax": (".parallax", "Parallax"),
+    "PaTHAttention": (".path_attn", "PaTHAttention"),
+    "Raven": (".raven", "Raven"),
+    "ReBasedLinearAttention": (".rebased", "ReBasedLinearAttention"),
+    "RodimusAttention": (".rodimus", "RodimusAttention"),
+    "SlidingWindowSharedKeyAttention": (".rodimus", "SlidingWindowSharedKeyAttention"),
+    "RWKV6Attention": (".rwkv6", "RWKV6Attention"),
+    "RWKV7Attention": (".rwkv7", "RWKV7Attention"),
+    "WallAttention": (".wall_attn", "WallAttention"),
+    "YOCOCrossAttention": (".yoco", "YOCOCrossAttention"),
+    "YOCOGatedRetention": (".yoco", "YOCOGatedRetention"),
+    "YOCOSharedKVBuilder": (".yoco", "YOCOSharedKVBuilder"),
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    spec = _EXPORTS.get(name)
+    if spec is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr = spec
+    value = getattr(importlib.import_module(module_name, __name__), attr)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/__init__.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy public exports so importing fla.ops.utils / fla.ops.gdn2 does not
+# eagerly compile every Triton kernel (needed on Windows + Triton 3.7).
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, str] = {
+    "chunk_abc": "fla.ops.abc",
+    "parallel_attn": "fla.ops.attn",
+    "fused_attnres": "fla.ops.attnres",
+    "fused_chunk_based": "fla.ops.based",
+    "parallel_based": "fla.ops.based",
+    "chunk_comba": "fla.ops.comba",
+    "fused_recurrent_comba": "fla.ops.comba",
+    "chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_recurrent_delta_rule": "fla.ops.delta_rule",
+    "parallel_forgetting_attn": "fla.ops.forgetting_attn",
+    "chunk_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "chunk_gdn": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gdn": "fla.ops.gated_delta_rule",
+    "chunk_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_gla": "fla.ops.gla",
+    "fused_chunk_gla": "fla.ops.gla",
+    "fused_recurrent_gla": "fla.ops.gla",
+    "chunk_gsa": "fla.ops.gsa",
+    "fused_recurrent_gsa": "fla.ops.gsa",
+    "fused_recurrent_hgrn": "fla.ops.hgrn",
+    "chunk_kda": "fla.ops.kda",
+    "fused_recurrent_kda": "fla.ops.kda",
+    "chunk_lightning_attn": "fla.ops.lightning_attn",
+    "fused_recurrent_lightning_attn": "fla.ops.lightning_attn",
+    "chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_recurrent_linear_attn": "fla.ops.linear_attn",
+    "chunk_log_linear_attn": "fla.ops.log_linear_attn",
+    "chunk_mesa_net": "fla.ops.mesa_net",
+    "parallel_nsa": "fla.ops.nsa",
+    "parallel_parallax": "fla.ops.parallax",
+    "parallel_path_attn": "fla.ops.path_attn",
+    "chunk_retention": "fla.ops.retention",
+    "fused_chunk_retention": "fla.ops.retention",
+    "fused_recurrent_retention": "fla.ops.retention",
+    "parallel_retention": "fla.ops.retention",
+    "chunk_rwkv6": "fla.ops.rwkv6",
+    "fused_recurrent_rwkv6": "fla.ops.rwkv6",
+    "chunk_rwkv7": "fla.ops.rwkv7",
+    "fused_recurrent_rwkv7": "fla.ops.rwkv7",
+    "chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_recurrent_simple_gla": "fla.ops.simple_gla",
+    "parallel_simple_gla": "fla.ops.simple_gla",
+    "parallel_wall_attn": "fla.ops.wall_attn",
+    "parallel_wall_attn_decode": "fla.ops.wall_attn",
+    "chunk_gdn2": "fla.ops.gdn2",
+    "fused_recurrent_gdn2": "fla.ops.gdn2",
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(importlib.import_module(module_name), name)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/simple_gla/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+from .chunk import chunk_simple_gla
+from .fused_chunk import fused_chunk_simple_gla
+from .fused_recurrent import fused_recurrent_simple_gla
+
+# Triton 3.7 on Windows can fail while decorating parallel kernels at import time.
+try:
+    from .parallel import parallel_simple_gla
+except Exception:  # noqa: BLE001
+    parallel_simple_gla = None
+
+__all__ = [
+    'chunk_simple_gla',
+    'fused_chunk_simple_gla',
+    'fused_recurrent_simple_gla',
+    'parallel_simple_gla',
+]