Kernels:

kernels-community
/

fla

Trusted publisher

kernels-bot commited on 4 days ago

Commit

e0d830e

verified ·

1 Parent(s): 2f06ade

Uploaded using `kernel-builder`.

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

README.md +1 -1
build/torch-cuda/__init__.py +2 -0
build/torch-cuda/_ops.py +1 -1
build/torch-cuda/{ops/backends/__init__.py → backends.py} +156 -114
build/torch-cuda/layers.py +10 -0
build/torch-cuda/metadata.json +264 -223
build/torch-cuda/metadata.json.sigstore +1 -1
build/torch-cuda/modules/__init__.py +86 -5
build/{torch-xpu/ops/cp → torch-cuda/modules/activations}/__init__.py +28 -21
build/torch-cuda/modules/{activations.py → activations/ops.py} +302 -359
build/torch-cuda/modules/backends/__init__.py +0 -17
build/torch-cuda/modules/backends/triton_ascend/__init__.py +0 -422
build/torch-cuda/modules/backends/triton_ascend/activations.py +0 -931
build/torch-cuda/modules/backends/triton_ascend/causal_conv1d.py +0 -1175
build/torch-cuda/modules/backends/triton_ascend/fused_cross_entropy.py +0 -469
build/torch-cuda/modules/backends/triton_ascend/fused_kl_div.py +0 -188
build/torch-cuda/modules/backends/triton_ascend/fused_linear_cross_entropy.py +0 -347
build/torch-cuda/modules/backends/triton_ascend/grpo.py +0 -266
build/torch-cuda/modules/backends/triton_ascend/layernorm.py +0 -355
build/torch-cuda/modules/backends/triton_ascend/rotary.py +0 -211
build/{torch-xpu/modules/conv/triton → torch-cuda/modules/causal_conv1d}/__init__.py +10 -3
build/torch-cuda/modules/causal_conv1d/backends/__init__.py +6 -0
build/torch-cuda/modules/{conv → causal_conv1d/backends}/cuda/__init__.py +2 -6
build/torch-cuda/modules/causal_conv1d/backends/cuda/cp.py +98 -0
build/{torch-xpu/modules/conv → torch-cuda/modules/causal_conv1d/backends}/cuda/ops.py +68 -87
build/torch-cuda/modules/causal_conv1d/backends/gluon.py +417 -0
build/torch-cuda/modules/{conv/cp/ops.py → causal_conv1d/cp.py} +149 -77
build/torch-cuda/modules/{conv/triton/kernels.py → causal_conv1d/ops.py} +766 -125
build/torch-cuda/modules/conv/causal_conv1d.py +0 -129
build/torch-cuda/modules/conv/cuda/ops.py +0 -233
build/torch-cuda/modules/conv/short_conv.py +0 -250
build/torch-cuda/modules/conv/triton/ops.py +0 -424
build/torch-cuda/modules/convolution.py +0 -42
build/torch-cuda/modules/feature_map.py +1 -2
build/torch-cuda/modules/fused_bitlinear.py +76 -559
build/torch-cuda/modules/fused_cross_entropy.py +0 -459
build/torch-cuda/modules/fused_cross_entropy/__init__.py +15 -0
build/torch-cuda/modules/fused_cross_entropy/gluon.py +402 -0
build/torch-cuda/modules/fused_cross_entropy/ops.py +538 -0
build/torch-cuda/modules/fused_kl_div/__init__.py +11 -0
build/torch-cuda/modules/{fused_kl_div.py → fused_kl_div/ops.py} +63 -113
build/torch-cuda/modules/fused_linear_cross_entropy/__init__.py +15 -0
build/torch-cuda/modules/{fused_linear_cross_entropy.py → fused_linear_cross_entropy/ops.py} +211 -443
build/torch-cuda/modules/fused_norm_gate.py +0 -1245
build/{torch-xpu/modules/conv/cp → torch-cuda/modules/grpo}/__init__.py +3 -5
build/torch-cuda/modules/{grpo.py → grpo/ops.py} +77 -117
build/torch-cuda/modules/l2warp.py +5 -6
build/torch-cuda/modules/{conv/long_conv.py → long_conv.py} +45 -63
build/{torch-xpu/modules → torch-cuda/modules/norm}/__init__.py +20 -22
build/{torch-xpu/modules → torch-cuda/modules/norm}/fused_norm_gate.py +197 -220

README.md CHANGED Viewed

@@ -11,7 +11,7 @@ This is the repository card of kernels-community/fla that has been pushed on the
 # make sure `kernels` is installed: `pip install -U kernels`
 from kernels import get_kernel
-kernel_module = get_kernel("kernels-community/fla", version=1)
 chunk_kda = kernel_module.chunk_kda
 chunk_kda(...)

 # make sure `kernels` is installed: `pip install -U kernels`
 from kernels import get_kernel
+kernel_module = get_kernel("kernels-community/fla", version=2)
 chunk_kda = kernel_module.chunk_kda
 chunk_kda(...)

build/torch-cuda/__init__.py CHANGED Viewed

@@ -1,3 +1,5 @@
 from . import layers
 from .ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
 from .ops.kda import chunk_kda, fused_recurrent_kda

+__version__ = "0.6.0"
 from . import layers
 from .ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
 from .ops.kda import chunk_kda, fused_recurrent_kda

build/torch-cuda/_ops.py CHANGED Viewed

@@ -25,7 +25,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "fla"
-    unique_id = "fe2bb66"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

 def _find_ops_name() -> str:
     kernel_name = "fla"
+    unique_id = "faa8674"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-cuda/{ops/backends/__init__.py → backends.py} RENAMED Viewed

@@ -5,26 +5,25 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-"""Generic backend dispatch system for FLA operations."""
 from __future__ import annotations
-import contextlib
 import logging
 import os
-import threading
 from collections.abc import Callable
 from functools import cache, wraps
 from typing import Any, ClassVar, TypeVar
 import torch
-from ...utils import find_spec_cached
 logger = logging.getLogger(__name__)
 F = TypeVar('F', bound=Callable)
 _DISPATCH_DISABLED = os.environ.get("FLA_DISABLE_BACKEND_DISPATCH") == "1"
 if _DISPATCH_DISABLED:
     logger.info("[FLA Backend] FLA_DISABLE_BACKEND_DISPATCH=1 — all dispatch bypassed")
@@ -36,6 +35,7 @@ class BaseBackend:
     Attributes:
         backend_type (str, Optional):
             Identifier for the backend type, used to distinguish different backend implementations.
             Default: `"base"`.
         package_name (str, Optional):
             Name of the external package required by the backend.
@@ -52,10 +52,20 @@ class BaseBackend:
     backend_type: ClassVar[str] = "base"
     package_name: ClassVar[str | None] = None
-    env_var: ClassVar[str | None] = None
     default_enable: ClassVar[bool] = True
-    # Lower number = higher priority, default is 5
     priority: ClassVar[int] = 5
     @classmethod
     def is_available(cls) -> bool:
@@ -63,161 +73,193 @@ class BaseBackend:
             return True
         return find_spec_cached(cls.package_name) is not None
-    @classmethod
-    def is_enabled(cls) -> bool:
-        if cls.env_var is None:
             return True
-        default_value = "1" if cls.default_enable else "0"
-        return os.environ.get(cls.env_var, default_value) != "0"
     @classmethod
     @cache
     def can_use(cls) -> bool:
-        return cls.is_available() and cls.is_enabled()
-    def verify(self, func_name: str, *args, **kwargs) -> tuple[bool, str | None]:
         """Check if backend can handle the function call."""
-        verifier_name = f"{func_name}_verifier"
-        verifier = getattr(self, verifier_name, None)
         if verifier is None:
             return True, None
         try:
             return verifier(*args, **kwargs)
-        except Exception as e:
-            return False, str(e)
-_OPERATION_BACKEND_MODULES: dict[str, str] = {
-    'modules': 'fla.modules.backends',
-}
-class BackendRegistry:
-    """Per-operation backend registry."""
-    _registries: ClassVar[dict[str, BackendRegistry]] = {}
-    _initialized: ClassVar[set[str]] = set()
-    _init_lock: ClassVar[threading.Lock] = threading.Lock()
-    def __init__(self, operation_name: str):
-        self.operation_name = operation_name
-        self._backends: dict[str, BaseBackend] = {}
-        self._active: BaseBackend | None = None
-        self._lock = threading.RLock()
-        self._logged: set[str] = set()
-        BackendRegistry._registries[operation_name] = self
-    def register(self, backend: BaseBackend) -> None:
-        """Register a backend."""
-        with self._lock:
-            self._backends[backend.backend_type] = backend
-            # Update active backend based on priority
-            self._update_active_backend()
-    def _get_sorted_backends(self) -> list[BaseBackend]:
-        """Get backends sorted by priority (lower number = higher priority).
-        Backends with the same priority are sorted by registration order.
-        """
-        return sorted(
-            self._backends.values(),
-            key=lambda b: (b.priority, list(self._backends.values()).index(b))
-        )
-    def _update_active_backend(self) -> None:
-        """Update active backend based on priority."""
-        for backend in self._get_sorted_backends():
-            if backend.can_use():
-                self._active = backend
-                return
-    def get_active(self) -> BaseBackend | None:
-        """Get active backend."""
-        return self._active
     @classmethod
-    def ensure_initialized(cls, operation: str) -> None:
-        """Lazy-load backends on first use."""
-        if operation in cls._initialized:
-            return
-        with cls._init_lock:
-            if operation in cls._initialized:
-                return
-            # Import backend module to trigger registration
-            module_path = _OPERATION_BACKEND_MODULES.get(
-                operation,
-                f'fla.ops.{operation}.backends',
-            )
-            with contextlib.suppress(ImportError):
-                __import__(module_path, fromlist=[''])
-            cls._initialized.add(operation)
-def dispatch(operation: str):
-    """Dispatch decorator with verifier support.
-    Iterates through all registered backends and selects the first one
-    that passes the verifier for the given function call.
-    """
-    def decorator(func: F) -> F:
         if _DISPATCH_DISABLED:
-            return func
-        func_name = func.__name__
-        @wraps(func)
         def wrapper(*args, **kwargs) -> Any:
-            # Lazy initialization of backends
-            BackendRegistry.ensure_initialized(operation)
-            registry = BackendRegistry._registries.get(operation)
-            if registry is None:
-                return func(*args, **kwargs)
-            # Iterate through all registered backends sorted by priority
-            # to find one that can handle this call
-            backends_list = registry._get_sorted_backends()
-            for be in backends_list:
-                # Avoid be.can_use(): its @cache wrapper breaks torch.compile tracing.
-                if not (be.is_available() and be.is_enabled()):
                     continue
-                can_use, reason = be.verify(func_name, *args, **kwargs)
-                if not can_use:
-                    fail_key = f"{operation}:{func_name}:{be.backend_type}:fail"
-                    if fail_key not in registry._logged:
-                        registry._logged.add(fail_key)
                         logger.info(
-                            f"[FLA Backend] {operation}.{func_name} -> {be.backend_type} "
                             f"rejected: {reason}"
                         )
                     continue
-                impl = getattr(be, func_name, None)
-                if impl is None:
                     continue
-                result = impl(*args, **kwargs)
-                log_key = f"{operation}:{func_name}:{be.backend_type}"
-                if log_key not in registry._logged:
-                    registry._logged.add(log_key)
-                    logger.info(f"[FLA Backend] {operation}.{func_name} -> {be.backend_type}")
                 return result
-            # No backend can handle this call, use default implementation
-            return func(*args, **kwargs)
-        # Dispatch performs runtime backend selection; keep it out of torch.compile graphs.
-        wrapper = torch.compiler.disable(wrapper)
-        return wrapper
     return decorator
-__all__ = ['BackendRegistry', 'BaseBackend', 'dispatch']

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+"""Backend registration and runtime selection."""
 from __future__ import annotations
+import importlib
 import logging
 import os
 from collections.abc import Callable
 from functools import cache, wraps
+from inspect import unwrap
 from typing import Any, ClassVar, TypeVar
 import torch
+from .utils import find_spec_cached, has_usable_nvcc
 logger = logging.getLogger(__name__)
 F = TypeVar('F', bound=Callable)
 _DISPATCH_DISABLED = os.environ.get("FLA_DISABLE_BACKEND_DISPATCH") == "1"
 if _DISPATCH_DISABLED:
     logger.info("[FLA Backend] FLA_DISABLE_BACKEND_DISPATCH=1 — all dispatch bypassed")
     Attributes:
         backend_type (str, Optional):
             Identifier for the backend type, used to distinguish different backend implementations.
+            `FLA_<BACKEND_TYPE>` enables all backends of this type, overriding individual switches set to `0`.
             Default: `"base"`.
         package_name (str, Optional):
             Name of the external package required by the backend.
     backend_type: ClassVar[str] = "base"
     package_name: ClassVar[str | None] = None
+    env_var: str | None = None
     default_enable: ClassVar[bool] = True
     priority: ClassVar[int] = 5
+    implementation: Callable | None = None
+    verifier: Callable | None = None
+    def __init__(self, *, env_var: str | None = None):
+        if env_var is not None:
+            self.env_var = env_var
+    def get_implementation(self, function_name: str) -> Callable | None:
+        if self.implementation is not None:
+            return self.implementation
+        return getattr(self, function_name, None)
     @classmethod
     def is_available(cls) -> bool:
             return True
         return find_spec_cached(cls.package_name) is not None
+    def is_enabled(self) -> bool:
+        if self.env_var is None:
+            return True
+        if os.environ.get(f"FLA_{self.backend_type.upper()}", "0") != "0":
             return True
+        default_value = "1" if self.default_enable else "0"
+        return os.environ.get(self.env_var, default_value) != "0"
     @classmethod
     @cache
     def can_use(cls) -> bool:
+        return cls.is_available() and cls().is_enabled()
+    def verify(self, function_name: str, *args, **kwargs) -> tuple[bool, str | None]:
         """Check if backend can handle the function call."""
+        verifier_name = f"{function_name}_verifier"
+        verifier = self.verifier or getattr(self, verifier_name, None)
         if verifier is None:
             return True, None
         try:
             return verifier(*args, **kwargs)
+        except Exception as error:
+            return False, str(error)
+class TritonAscendBackend(BaseBackend):
+    """Shared availability and selection policy for Ascend implementations."""
+    backend_type = 'triton_ascend'
+    priority = 0
+    @classmethod
+    def is_available(cls) -> bool:
+        from .utils import IS_NPU
+        return IS_NPU
+class TileLangBackend(BaseBackend):
+    """Shared availability and selection policy for TileLang implementations."""
+    backend_type = 'tilelang'
+    package_name = 'tilelang'
+    env_var = 'FLA_TILELANG'
+    @classmethod
+    def is_available(cls) -> bool:
+        return super().is_available() and has_usable_nvcc()
+class GluonBackend(BaseBackend):
+    """NVIDIA GPU backend using Gluon kernels."""
+    backend_type = 'gluon'
+    package_name = 'triton.experimental.gluon'
+    env_var = 'FLA_GLUON'
+    default_enable = False
     @classmethod
+    def is_available(cls) -> bool:
+        from .utils import IS_NVIDIA
+        return IS_NVIDIA and super().is_available()
+class BackendRegistry:
+    """Ordered backend candidates for an operation or entry point."""
+    def __init__(self, operation_name: str):
+        self.operation_name = operation_name
+        self._backends: dict[str, BaseBackend] = {}
+        self._logged: set[str] = set()
+    def register(self, backend: BaseBackend) -> None:
+        """Register or replace a backend by its type."""
+        self._backends[backend.backend_type] = backend
+    def _get_sorted_backends(self) -> list[BaseBackend]:
+        """Lower priorities come first; ties retain registration order."""
+        return sorted(self._backends.values(), key=lambda backend: backend.priority)
+    def get_active(self) -> BaseBackend | None:
+        """Return the first available and enabled backend, before call verification."""
+        for backend in self._get_sorted_backends():
+            if backend.is_available() and backend.is_enabled():
+                return backend
+        return None
+    def dispatch(self, default_function: F) -> F:
+        """Select the first eligible backend, falling back to the decorated function."""
         if _DISPATCH_DISABLED:
+            return default_function
+        function_name = default_function.__name__
+        @wraps(default_function)
         def wrapper(*args, **kwargs) -> Any:
+            for backend in self._get_sorted_backends():
+                # avoid backend.can_use(): its @cache wrapper breaks torch.compile tracing.
+                if not (backend.is_available() and backend.is_enabled()):
                     continue
+                accepted, reason = backend.verify(function_name, *args, **kwargs)
+                if not accepted:
+                    rejection_key = f"{self.operation_name}:{function_name}:{backend.backend_type}:fail"
+                    if rejection_key not in self._logged:
+                        self._logged.add(rejection_key)
                         logger.info(
+                            f"[FLA Backend] {self.operation_name}.{function_name} -> {backend.backend_type} "
                             f"rejected: {reason}"
                         )
                     continue
+                implementation = backend.get_implementation(function_name)
+                if implementation is None:
                     continue
+                result = implementation(*args, **kwargs)
+                log_key = f"{self.operation_name}:{function_name}:{backend.backend_type}"
+                if log_key not in self._logged:
+                    self._logged.add(log_key)
+                    logger.info(f"[FLA Backend] {self.operation_name}.{function_name} -> {backend.backend_type}")
                 return result
+            return default_function(*args, **kwargs)
+        # runtime backend selection must stay outside torch.compile graphs.
+        return torch.compiler.disable(wrapper)
+_operation_registries: dict[str, BackendRegistry] = {}
+_function_registries: dict[Callable, BackendRegistry] = {}
+def _get_operation_registry(operation: str) -> BackendRegistry:
+    if operation not in _operation_registries:
+        _operation_registries[operation] = BackendRegistry(operation)
+    return _operation_registries[operation]
+def _load_operation_registry(operation: str) -> BackendRegistry:
+    # an existing registry does not guarantee that all of the owner's backends are loaded.
+    importlib.import_module(f'.ops.{operation}.backends', __package__)
+    return _get_operation_registry(operation=operation)
+def register(
+    entry_point: Callable | str,
+    *,
+    backend: type[BaseBackend] | None = None,
+    verifier: Callable | None = None,
+    env_var: str | None = None,
+) -> Callable[[F], F]:
+    """Register an implementation or adapter for a dispatched function or operation name."""
+    def decorator(implementation: F) -> F:
+        if _DISPATCH_DISABLED and not isinstance(entry_point, str):
+            return implementation
+        if isinstance(implementation, type):
+            registered_backend = implementation()
+            if backend is not None and backend.backend_type != registered_backend.backend_type:
+                raise ValueError('The registration backend must match the adapter backend_type')
+        else:
+            registered_backend = backend() if backend is not None else BaseBackend()
+            registered_backend.implementation = implementation
+        registered_backend.verifier = verifier
+        if env_var is not None:
+            registered_backend.env_var = env_var
+        if isinstance(entry_point, str):
+            registry = _get_operation_registry(operation=entry_point)
+        else:
+            registry = _function_registries[unwrap(entry_point)]
+        registry.register(registered_backend)
+        return implementation
     return decorator
+def dispatch(entry_point: F | str) -> F | Callable[[F], F]:
+    """Allow backends to register replacements while retaining the original fallback."""
+    if isinstance(entry_point, str):
+        return _load_operation_registry(operation=entry_point).dispatch
+    if _DISPATCH_DISABLED:
+        return entry_point
+    registry = BackendRegistry(entry_point.__module__)
+    _function_registries[unwrap(entry_point)] = registry
+    return registry.dispatch(entry_point)
+__all__ = ['BackendRegistry', 'BaseBackend', 'dispatch', 'register']

build/torch-cuda/layers.py CHANGED Viewed

@@ -7,6 +7,8 @@ from .ops.kda import chunk_kda, fused_recurrent_kda
 class FusedRMSNormGated(nn.Module):
     def forward(self, hidden_states, gate=None):
         return rms_norm_gated(
             hidden_states,
@@ -22,6 +24,8 @@ class FusedRMSNormGated(nn.Module):
 class chunk_kimi_delta_attention(nn.Module):
     def forward(
         self,
         query,
@@ -54,6 +58,8 @@ class chunk_kimi_delta_attention(nn.Module):
 class recurrent_kimi_delta_attention(nn.Module):
     def forward(
         self, query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False, **kwargs,
     ):
@@ -75,6 +81,8 @@ class recurrent_kimi_delta_attention(nn.Module):
 class chunk_gated_delta_rule(nn.Module):
     def forward(
         self,
         query,
@@ -107,6 +115,8 @@ class chunk_gated_delta_rule(nn.Module):
 class recurrent_gated_delta_rule(nn.Module):
     def forward(
         self, query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False, **kwargs,
     ):

 class FusedRMSNormGated(nn.Module):
+    can_torch_compile = True
     def forward(self, hidden_states, gate=None):
         return rms_norm_gated(
             hidden_states,
 class chunk_kimi_delta_attention(nn.Module):
+    can_torch_compile = True
     def forward(
         self,
         query,
 class recurrent_kimi_delta_attention(nn.Module):
+    can_torch_compile = True
     def forward(
         self, query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False, **kwargs,
     ):
 class chunk_gated_delta_rule(nn.Module):
+    can_torch_compile = True
     def forward(
         self,
         query,
 class recurrent_gated_delta_rule(nn.Module):
+    can_torch_compile = True
     def forward(
         self, query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False, **kwargs,
     ):

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,7 +1,7 @@
 {
   "name": "fla",
-  "id": "_fla_cuda_fe2bb66",
-  "version": 1,
   "kernels-minver": "0.14.0",
   "license": "MIT",
   "python-depends": [
@@ -14,297 +14,338 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "bVRNYUhThvULeUaovvQFZCJ9IZZsC+0kU/HbnXdYT1A=",
-      "_ops.py": "pXgvH5hf/dXVr6/amV+mRqgiBUGncbV3gbMubhmlOW4=",
-      "layers.py": "/39XVQW4d4GHeNROqT67ujKbItdteLBtm0PT6zq41YQ=",
-      "modules/__init__.py": "6ZjjRMt6Rf+2OegXDC8w58RD1IW3SBK+dxBoZa4Uohc=",
-      "modules/activations.py": "Lp7xLUf7r785nt93NFhqoSQM+G7AK2e6kviE2/0TRJs=",
-      "modules/backends/__init__.py": "l0tLqWyPFwwDRTdnHZYBzoux6J/LOFveJ53KREAKKDk=",
-      "modules/backends/triton_ascend/__init__.py": "vbU1UZwQDjVfF75INGYuWlIhuUTzx2TuZHIvme/QSMw=",
-      "modules/backends/triton_ascend/activations.py": "eqE/iXLZZEan7XHSa9ziz7mWNYwxTC5kIcJiwDwRr/Y=",
-      "modules/backends/triton_ascend/causal_conv1d.py": "VRxcTWPBB221dgd0+eNHEzCi9Echi6B2XOINtNOjHsw=",
-      "modules/backends/triton_ascend/fused_cross_entropy.py": "RITRIWoSOKhG8Mz7odkko27gESEd8M48ICc1IXRSVPg=",
-      "modules/backends/triton_ascend/fused_kl_div.py": "jb7XzKT6AKmXl8NxjZ4/SNofjdD3ioi4Df/l4KsIAtI=",
-      "modules/backends/triton_ascend/fused_linear_cross_entropy.py": "Z8/wQfPzJTg9ZICxdsABldRPsMa+mUZUTHhCJS2OeQc=",
-      "modules/backends/triton_ascend/grpo.py": "Jy+RkxbJ4RgJsEbk0xNgnw/D3Gj/P+6xDaMfx3I7DPE=",
-      "modules/backends/triton_ascend/layernorm.py": "7uKqeIraFdvcAVR0g6CLw+NJDCX8cLRMjNGHoFx1AnM=",
-      "modules/backends/triton_ascend/rotary.py": "+CmvapJJ3AOlwHbRxxSEIeVDjmlFPRno+5aqdfAHznI=",
-      "modules/conv/__init__.py": "WZKA4OJWHEN0bg2Yl2i2dvfkSp9GZKaD65abY8pSKbw=",
-      "modules/conv/causal_conv1d.py": "sHN5nlzDzSQwgo40OPt9lzcLRzZZnPCFvS/pmfK8L9k=",
-      "modules/conv/cp/__init__.py": "BKKfEjTjwyEAwDu66avCFolqnmrv9qHHNvxjfOHi3fY=",
-      "modules/conv/cp/ops.py": "pv7PGhmkfBsFt4tY8Hd0UAHzUPqp6DhgjRJmUujWitU=",
-      "modules/conv/cuda/__init__.py": "q35Kd12ITYtHp5JtaAhQJ6/YCEY95RWWveKoFPJ/VCU=",
-      "modules/conv/cuda/ops.py": "YE2httgrnIB1iSE5Uw9h1wtnq0UtybPuouF3z2reUb0=",
-      "modules/conv/long_conv.py": "h+wI2DdlBSQIZC335mwgBsjxGbByGaklb6Nv4WcQMMg=",
-      "modules/conv/short_conv.py": "DdAbuCpegRP32xEECOQQ4Ty3gWAOiFUzJLCD2mxcU1w=",
-      "modules/conv/triton/__init__.py": "fDLG2FdP68jqqsBxNdKifAIEjmvGzXNrLm/xDMrfbvg=",
-      "modules/conv/triton/kernels.py": "EK+wcpRMAGXKGmFCrZzxEeuZJmWvzUqxdTBYLwrpFT8=",
-      "modules/conv/triton/ops.py": "jIxuam5Ch5U3ErbT8D6Zn/5FqBoTq15kcJ9BGMD8exU=",
-      "modules/convolution.py": "WREsoV3gCce7rMEdpuEsyBgvmfbhk2IFdpUrr4TMtBw=",
-      "modules/feature_map.py": "sfoRvsdgvhtWkKuMYE/f1c7eyINwEaZ9TfItktWD+RE=",
-      "modules/fused_bitlinear.py": "lWTe4Ym60A+rHZTdQ6cZqa1Cje5vG/xMWvzLSAzLpQc=",
-      "modules/fused_cross_entropy.py": "si+ZNK/vu/TERnt/v2/Tm+ZBfn3poUoPDDTuJmrEuYM=",
-      "modules/fused_kl_div.py": "IDl78EMB9KXWmbV2cflnbXmcQpzdoqkwlEj4tvAPsbA=",
-      "modules/fused_linear_cross_entropy.py": "SruIxSe0Hhlf7n0upOldvn9HwOfL6XQ42N6qGZC6aOA=",
-      "modules/fused_norm_gate.py": "LAOhNgUSEiIU3Y5Ya5EHRjMAxHT0StODm04AUcL4Nys=",
-      "modules/grpo.py": "JPZnSGESIMmHDVkfYDSUzbm1qb2vY6dmDmFvSMI53RU=",
-      "modules/l2norm.py": "LchUl4mEXyx2cgHPuNgUhkorbNeuN+P7SfJyGXX0OyI=",
-      "modules/l2warp.py": "5oJlj6JRXBGe5ZANNnUA8GLIjlzXCnddtz2lwgKn6WA=",
-      "modules/layernorm.py": "1QhCZYF7pFRmZTxCWC55oG1uB6Lak5RXYca8MF5q5SA=",
-      "modules/layernorm_gated.py": "s/c/gRsZ3O6jXwY3b9owPaN1QAoOvsSevWuW2EXuJjE=",
       "modules/mlp.py": "q09Zw5Bttqq27hnlybbQ9ncr0jSw1UCrJH2Ewh9roV4=",
-      "modules/parallel.py": "y3ZBLQp3FeDfBlHYmSerllqbfE3RgzX4k26lA15DPeg=",
-      "modules/rotary.py": "qVsb20GcFf9RolZ3dChWzEi5S9rc5UaMdC3dlrhNUj4=",
-      "modules/token_shift.py": "IWlvUPZ/Zm4IKW0dWSsx4p4r+ovbpPHQhk5Dp09P5Wg=",
-      "modules/token_shift_cp.py": "Zvaa6d/GiuljKqNgLdkjX0vysktpIVV8KMSTOjdG3Hw=",
-      "ops/__init__.py": "PsKzHYHoje01DfrVdltffiZFdS5CTVps5xl9dXr2RHo=",
       "ops/abc/__init__.py": "WwLiIkYi3xJ5RiG/bOCNNRaeeqmMuIPT6n1g23qfpDg=",
-      "ops/abc/chunk.py": "sqvLGP2r+frjpL2qWK8GHCoWWPU5Df8IT5y20oaTZTE=",
       "ops/abc/naive.py": "AZ3uz9eDvycW3gTwZuv9J4k86KndfQMZiLBFyPmzsLw=",
       "ops/attn/__init__.py": "gU/67mJonbeXpzfxgi8K7HIG2a8Xpjd0gOuCk/Pu9EI=",
-      "ops/attn/backends/__init__.py": "3xcuWH+iLBPRb3g4ZkyXBXm8wJFxQK4GHdUpmK8Nh88=",
-      "ops/attn/decoding.py": "ScKiTKcbl4kcnYDNz6jXd2a1qB5ymDcz8YVkotRnO9Y=",
-      "ops/attn/naive.py": "nBsWqmBWOMQdaQYvdFejDLoX0Pei9W8oMPrHZoCCin0=",
-      "ops/attn/parallel.py": "+Vkkjy4OIwR+7SXBf6fCJqrPg740qc6a1+21RSNt9Cs=",
       "ops/attnres/__init__.py": "/ljBk1VNFjDa9XABPXqIiQ05vUdzRV6Dbx/KkQihrWs=",
-      "ops/attnres/fused.py": "Q8RKiwlUtFZkHELCLV3wtQGf1cVDPoDNMK1waEwbxFQ=",
       "ops/attnres/naive.py": "nV4WPz/vfRwfRsIzF7pKV95fEx2dQP4BrrPoQCh0a0M=",
-      "ops/backends/__init__.py": "3DU4e6t7dmPuZLlHjPnt4q/MAtDXtaMmAcSY4ycyytY=",
       "ops/based/__init__.py": "qdBLUigpgjQaIN3HTtq57uxoqAaBnQkZbFiEIf18Rvk=",
-      "ops/based/fused_chunk.py": "Evgk49vwK6VxWWmw3KdwhiUV5N8RF/PcxZP8J1dlYR8=",
       "ops/based/naive.py": "YWIFSLeslp2ruyQDz/uCr5V4ao2cyRyeQH1wiVLMB8U=",
-      "ops/based/parallel.py": "7hNqqI2UIeJNebDTF7RwW+Q9ST15DwMxJugiaVO9s/k=",
       "ops/comba/__init__.py": "FBKonbvPPFHY7an2Auh1sXVJLbcgwX8aAD5fxi2oecY=",
-      "ops/comba/chunk.py": "o+3soCUQ27uZXe+YOZpYF1/aiLFEyTdbb+bmetY6O2c=",
-      "ops/comba/fused_recurrent.py": "ujg1r9D2Jx+/sgli1msfkXoiCZb2gX9dL/OpOca3kGs=",
       "ops/comba/naive.py": "QH967c5rxlSSDs6J6jRtLjZ3iLiwgAYgY9+0m9rgD18=",
-      "ops/comba/utils.py": "rcW3h8gjoEzSoLgY1HkrIwsn9XKEYtz274CGqr8wWzU=",
-      "ops/comba/wy_fast.py": "EsJfyG3HDGY9rgJ97cn2SAq1DeqnXdgj5mzdOlWtz2I=",
-      "ops/common/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=",
-      "ops/common/backends/__init__.py": "aEhY8S2W0I6qbjedaCYjqLt4VaNumWM+bjKtPEJSAN0=",
-      "ops/common/backends/intracard.py": "yXACesLmdhG1PwAFJuquz4F+hqSNtFvXwZ51YeEN8FM=",
-      "ops/common/backends/tilelang/__init__.py": "2Zz/LkJfy8oOvqUcen6iQDMcf2prznZpYhHAc5QYW3E=",
-      "ops/common/backends/tilelang/chunk_bwd.py": "5zXyBzqIXRgoZdZ8zjLNEIz3jBksHpaVxnqUk66ckjQ=",
-      "ops/common/backends/tilelang/parallel_attn_bwd.py": "w+KjVu0voRjcFhmwWfIEYaSwmlk17oIu9EKRvjNIAro=",
-      "ops/common/backends/tilelang/parallel_attn_fwd.py": "7INMyGs1d3H8BKcjRnMvu/NHMm9VeNFTETR5z7ZO15c=",
-      "ops/common/chunk_delta_h.py": "QmkSsbbyOHai9cVaMB8O/8ep4NxiM6ieUVKmja6USw8=",
-      "ops/common/chunk_h.py": "imyz71p9AD6nOtStpN+zzgE0yVbGnq4DO3FitrO3hvc=",
-      "ops/common/chunk_h_parallel.py": "YGNR5a+eC7UShlVAvQF3fSLK+HXoP2Nc06H3p/PyzVg=",
-      "ops/common/chunk_h_split.py": "zYoWf1Li3Q7MFPusp5BRf1F+OuGzxxTeiZ3hdunYnCY=",
-      "ops/common/chunk_o.py": "1yUHQ+B2J1PQpxTFEa8R6JxpyV1k/TNXiXkOvFRfQAo=",
-      "ops/common/chunk_scaled_dot_kkt.py": "MyoKTPKUUZlOzdC4FdP4etL6pPBGzaNL5DJQyH3603E=",
-      "ops/common/fused_chunk.py": "nKIxBC8UD00KuY7WGFysC8P4ELNLEJhHmT60VcqyHbc=",
-      "ops/common/fused_recurrent.py": "RnQdVD59iOcRa8sSK4XcQT59RhZJr42YplDaQrDt97o=",
-      "ops/common/gate.py": "ExEV4w/h/y44Qb6L1cwOFEsMutGjs/t338KimPVjeHM=",
-      "ops/common/intracard_cp.py": "R+wNrYC2m0gUn0jr5CpsxpCIVRCRkRnFXjBus9XbDE8=",
       "ops/cp/__init__.py": "uud2WkodJvH5U5DjBa5oSLJtZ1G8O4D4XEuspXyC9lA=",
-      "ops/cp/chunk_delta_h.py": "+pPn7+hNImlTrCZLHNpaJKabA+OaoNXofYbX2F4fTD0=",
       "ops/cp/comm.py": "lqyQDwmK19vowTtGJGO+6SEn/neFFE0G8xFvlxLOmb0=",
-      "ops/cp/context.py": "N/p87Th0z6m2cGhP6icCN7rEg3QPVua8ymP8oDD4qy0=",
       "ops/delta_rule/__init__.py": "1gIkBajhtbVwuJOgPQ+FYYdq4tYgNgib+veOnont+N4=",
-      "ops/delta_rule/chunk.py": "r1b6g3/4OO1oDZ8n/MMyzagYHs0c5FPKU8//7KskU8w=",
       "ops/delta_rule/fused_chunk.py": "2NiUBRAl98YgRVGr2uzYJV22pQW3WNOe6euTX1Mdqyw=",
-      "ops/delta_rule/fused_recurrent.py": "aerY1sPJ0y6IknruBVOwbu3v/pLz2PjP3aNEzsArp5M=",
       "ops/delta_rule/naive.py": "IK4DL04IAZtCmMyfzyw5fYcfqsD9HxgCQ9n2makZKWg=",
-      "ops/delta_rule/parallel.py": "GMhpAZoEDuzyIjK0Qcpu+ku528WnLgDUujcsoRnXU6o=",
-      "ops/delta_rule/wy_fast.py": "O0UuKiOw3mImzRQHFLKHifSewFAoqRJ9i9r0puiX9fw=",
       "ops/deltaformer/__init__.py": "51ZmNFoHy3P0lVn/2EHKu1QS2BMC6h+hnAwSIr9szYg=",
       "ops/deltaformer/invcum.py": "ren181AxlHBv2MoZIHmSGMmJ0sW6HDb3t8fqG7gDELU=",
-      "ops/deltaformer/naive.py": "8b88rviKkDZpW+F9GwgEkVkwuSioECwgRJX//CVksnI=",
-      "ops/deltaformer/parallel.py": "sDHrnr2fM4+0T2cyAJoI2VP9Lan3NmaaDxk9JLtyLWk=",
       "ops/dsa/__init__.py": "ZNyws889LPBiY8gwB5zTs2vQwJYV8bX5jg6Cldkkq9s=",
-      "ops/dsa/naive.py": "lKgrdKOoE/qCdF4DtrI0uEkRdB28kUYbtn8hO63fzyo=",
       "ops/forgetting_attn/__init__.py": "A6ozg9zdEsz4lgUAXqEIUY8CeNp+5IZqYKOpChpj/A4=",
-      "ops/forgetting_attn/naive.py": "RDPXw01ZnhTMV125OFvEVeDDMht3BYj5gQI4oaT+T1M=",
-      "ops/forgetting_attn/parallel.py": "fcKF7Q84XXmeeKI2Tpv5XIfmxFovuXlmBov7Ie8OjKg=",
       "ops/gated_delta_product/__init__.py": "0MAzYdxP0ODZjN4WNQ+3HLwigokyfMv15rXUpTX+qnw=",
-      "ops/gated_delta_product/chunk.py": "QNHYG+1YKzot3cQ2oWvzuv7nDoZjllEXEomCk11WUJM=",
-      "ops/gated_delta_product/chunk_deltaproduct_h.py": "5IvD8ZZxaN216Wj8dop1cwJE9P5rvZ+Si98qB6kC5DQ=",
-      "ops/gated_delta_product/chunk_deltaproduct_o.py": "nvMmsV+E9AG4Z9HA6UV69DdGMUYH2XcL8NAAubUtIHc=",
       "ops/gated_delta_product/chunk_ref.py": "R1dHdPy4rU+lTQQKodAn9h2xPLxmwI2z0T6pRvh8Llo=",
-      "ops/gated_delta_product/naive.py": "Q907vtcavFpalFlttThkWSb/SaTa4vqAG+I/nWvKDno=",
       "ops/gated_delta_rule/__init__.py": "Xt0hhHr3CECyZDmBIZakdufk3bvvIa/V/iBmpkYwpnw=",
-      "ops/gated_delta_rule/chunk.py": "DFEZWmp016p+65pnDynAMVOY+NsUTGOfVcXsAhQr1i4=",
-      "ops/gated_delta_rule/chunk_fwd.py": "NCMKqzu9uM2yAGhvVqENRtoVLwoxVo/KaVChkvIJIDE=",
-      "ops/gated_delta_rule/fused_recurrent.py": "M7x3/U1Xnf5mndUc+lvGpzxu4KDNc2Sfevv+RVZZFNk=",
-      "ops/gated_delta_rule/gate.py": "H9JLPKjgZmorhyV+1IJFSMIa6x9admDn+EXXlwg167g=",
       "ops/gated_delta_rule/naive.py": "0c8XmSNJ/T6Ur5mbIuPTqBvkotGIHOW3CjRXJXFm4Ms=",
-      "ops/gated_delta_rule/wy_fast.py": "kFbW7BX1I4DBBBhfV1yXGQR4Jb9tGM6yILNmTXNOhxo=",
       "ops/gated_oja_rule/__init__.py": "oShC/cL0jMy+xWR7YBV9QxkAdxt1xasXDCypGAU6Hgo=",
       "ops/gated_oja_rule/chunk.py": "lq+5BXEaJLfL4dqqcJb4TmtyF9OrP9Wh+JpRfpdzd2s=",
-      "ops/gated_oja_rule/chunk_h.py": "kr6KBcPbFRrIQFBoQjmt1wsyvO1pyLeGW5v/k+NI87Q=",
-      "ops/gated_oja_rule/chunk_kkt.py": "N/sU3E2i5+jQBXnFuWyaUAiQkAgWUxNjO9lxdlpEkDA=",
-      "ops/gated_oja_rule/chunk_o.py": "4tPVh8KWj+D2onfXbIhnhaGWq35az5uJUwxTUbbRll0=",
-      "ops/gated_oja_rule/fused_recurrent.py": "yPfAcIXWBlr50lvK59bbF2qkhKWt40eT+EXx6exojK0=",
-      "ops/gated_oja_rule/wy_fast.py": "mS9GSN2RVUbEMkz2o1Hfi6j0nOhXtjLLSIccXkwmW3c=",
       "ops/gdn2/__init__.py": "xiEsgvVEfHhKcurQqL7ZKGGzEBW2BtGAylE9vxoNmAk=",
-      "ops/gdn2/chunk.py": "zgJK/Mwhn1fA9PpvagLdy6yFuxo2c1Z8DmUY5FJp2kw=",
-      "ops/gdn2/chunk_bwd.py": "Ay8kHSsnyD1wxGE1YcEfV3d4L3yae/8bDBMReKPbVBw=",
-      "ops/gdn2/chunk_fwd.py": "7Z0d9SnKSby4yCAeXx5kKx6HIdYfaXl0ZhYDDAniClU=",
-      "ops/gdn2/chunk_intra.py": "UjnmvfEg9EB41jGA9YHcZT9R2oHpBQI4PCq1Dwj1ZiQ=",
-      "ops/gdn2/chunk_intra_token_parallel.py": "U1lf9FZ6mfW2iIm06YJ5qPX8IQeT4U/G+J7RlZoWutw=",
-      "ops/gdn2/fused_recurrent.py": "3QsJl80MRQYrQXVN9lmgw4cVD1hLHrQXg6xcIgU5Bz8=",
       "ops/gdn2/naive.py": "amWAVAHrMA34W2KXYRoN0VjdGDQjzyvNXg8Ozd4+nno=",
-      "ops/gdn2/wy_fast.py": "DNH9FhRaZl4nqZ/YrxH2SqCgZdC77bkJ24sFIemnkRM=",
       "ops/generalized_delta_rule/__init__.py": "cn6rKWEDrxWw428Dk6EEy4Y7nyWHpc8rtR+EbRO9ndA=",
       "ops/generalized_delta_rule/dplr/__init__.py": "F8TbaiZeRKjbKhlI4XYopNd4bfzV10feNO06i32h2Vo=",
-      "ops/generalized_delta_rule/dplr/chunk.py": "qeR4srXoGT8NjNx93IQEyPGq3sqsK78CdOovbRwbntU=",
-      "ops/generalized_delta_rule/dplr/chunk_A_bwd.py": "MDoA0HOWDti16rAh8Q4D4RJonJkmIN04dCwUsHOM+sk=",
-      "ops/generalized_delta_rule/dplr/chunk_A_fwd.py": "wQSOE/PoqvjP6ugh3wSmbsst0ECSclxJEl5x/4NQ0KI=",
-      "ops/generalized_delta_rule/dplr/chunk_h_bwd.py": "hAISrlVF9QETkcoDAgEKXRY0w9DSDKdohqTXfkg421E=",
-      "ops/generalized_delta_rule/dplr/chunk_h_fwd.py": "vWtEs8L4DXX52UF5qt6Vuu/pOsHzKOFOIFtyezheWPY=",
-      "ops/generalized_delta_rule/dplr/chunk_o_bwd.py": "CuHt6fYI2g/6m0k2mj7JSLhCQ+zcKrMUYc2LOmKuUUs=",
-      "ops/generalized_delta_rule/dplr/chunk_o_fwd.py": "KOF8wWIEMIyhDfSL44bucPPsCp6u3AdYPlRDKK10KyQ=",
-      "ops/generalized_delta_rule/dplr/fused_recurrent.py": "ABaz+l4n77U9FrE3bdsGx/5jexAnvJ1cJcI/XJkeIac=",
       "ops/generalized_delta_rule/dplr/naive.py": "4flLSZUtwaq5S/PH+fu04rzOM2+USdRX6Ofte4lanE8=",
-      "ops/generalized_delta_rule/dplr/wy_fast_bwd.py": "HiMBn/U+a0vSCOVmn29UYLylI5x26qqualdl4E6ipd0=",
-      "ops/generalized_delta_rule/dplr/wy_fast_fwd.py": "R+CamNa7b49LfvrNZQs4roi+TReWcLGvBcoT3DbhoLM=",
       "ops/generalized_delta_rule/iplr/__init__.py": "6Vg8MTc+Byrnmx3HJz2yySi9nk2kyKCSg7aL4CMJ+GA=",
-      "ops/generalized_delta_rule/iplr/chunk.py": "BsCf/fiYn3wpS0eDPrR+mCrqsAihXD5v/QD2QkA5ATc=",
-      "ops/generalized_delta_rule/iplr/fused_recurrent.py": "nNo+ww/QmydOqBLnpmxXHAbANGDPTeWGjFI7nA9Vrw0=",
       "ops/generalized_delta_rule/iplr/naive.py": "YotccLOueXwBcQys4ICmFAysSrX8fj9eRasXIpmpRMo=",
-      "ops/generalized_delta_rule/iplr/wy_fast.py": "SoTcF2asZEIAnZS1qR9PXHGkHwnBspjuUsz6RpuqNiU=",
       "ops/gla/__init__.py": "PTUqqQlC4cXrhRhG/KRx+7X3m9WJKxGyzrIB/WJEXSo=",
-      "ops/gla/chunk.py": "B7reIP0UzTvYgK8Ds50r0uABlaucJWm70JVKKdYe1L0=",
       "ops/gla/fused_chunk.py": "OjzKcjSQujVQ+Q+d1Lmg4tsTzZaK6xWs8AGuplRn5zo=",
-      "ops/gla/fused_recurrent.py": "ZIzMstFj8WAXagDzTlZo3hFZy0k+egHNg8Qr61sb4vw=",
       "ops/gla/naive.py": "rEuH5bEAieya1CTvtZbebVFo2J5LdOX1ztL1WiDkZB0=",
       "ops/gsa/__init__.py": "9uygbV6orNtaTkZl9Hz+aAp7EvT6RuffawfeFtuUNEE=",
-      "ops/gsa/chunk.py": "SsfwmTOpKRG2Fja/lArJkw1teaomk/Vgczjr2/Oa+l0=",
-      "ops/gsa/fused_recurrent.py": "K7qBWOmh0uL9CO7D3pSlrxKFHy1fobTUyTFV5p7okXg=",
-      "ops/gsa/naive.py": "rezNszsU38KD1iHBW+ErV/Z01T5LIPadEjJnWiXz9LI=",
       "ops/hgrn/__init__.py": "69DMv7FQYvobQTXqOFtu7muAt1k1k9nvsFXcMAmjQmg=",
-      "ops/hgrn/chunk.py": "Br3GeHkw1EcOF/4X2BOvUTclhfLDZwkthIshoHTMMqg=",
-      "ops/hgrn/fused_recurrent.py": "KvO/5l/3wnJiT0jT0UA9IYxQhGHJy8jC7lYaOOyLM1w=",
       "ops/hgrn/naive.py": "XKo+5zJ+G4elkgK5ARsoXm5WYkG1QidTKNIuPT4k588=",
-      "ops/kda/__init__.py": "JFZPAQH4eiYFbpBh7Xccrw2PbZoAtNve1wHqdSW0Wss=",
-      "ops/kda/backends/__init__.py": "pk1xLku49AblKggFmQjmGlr1oNqUgdzxggzzgx8yeZk=",
-      "ops/kda/backends/flashkda.py": "g63Eq+K84JInsin5NVb428RIRPhL3wH+6Y4a+OHSFqk=",
-      "ops/kda/backends/tilelang/__init__.py": "Y+1//U4pzuPg2mbHoCymDfs186j482cI85ZOPxi88E0=",
-      "ops/kda/backends/tilelang/chunk_bwd_dqkg.py": "Fk1eT3M/KPFvfVRrJQ6C0158h5GaR72YK3tgSxkKoUo=",
-      "ops/kda/chunk.py": "LnoIH4MI8gHR8613fOKLqJGVNQkbtzjEJ0R7qX8/kQE=",
-      "ops/kda/chunk_bwd.py": "AahkP3AepxWsoyGCZHlwUGc7rGIHJh68Im5r/Zs3MK0=",
-      "ops/kda/chunk_fwd.py": "wUTWWQVWLRGeDj9e7dO03+0IosaRuWAfYqsl2KBy0ZA=",
-      "ops/kda/chunk_intra.py": "Ao7C9D/HzYJkiFxIL5s/LIlVE2YNgpL9yASJl5giPOs=",
-      "ops/kda/chunk_intra_token_parallel.py": "ENjjhDAR0cK2mqoU4q7QpBMPHmICgoLO06e4dEnd5hY=",
-      "ops/kda/fused_recurrent.py": "r6ZJMiLVmCRLpbYXS25mIX58KU2yKtDJqGwwHTXY0MY=",
-      "ops/kda/gate.py": "5056ezuomW2CNvPE04MraGJtf0gHvGqicxin1H0DICg=",
       "ops/kda/naive.py": "YKMihdS2cGj/YztIu+irMQKAZtJPANJ+EhmaiPxz8BY=",
-      "ops/kda/wy_fast.py": "9pMMqS5E/l0mFV/Tyap1QKY39CB5Iy94/Yn4s9WzjUw=",
       "ops/lightning_attn/__init__.py": "FPVVYNLjfO3GJQYIp+4HCQ4MdJtF8UZx5yktWDIatFA=",
-      "ops/lightning_attn/chunk.py": "kFjfxpUOVj61TnH0vXWZhDnw318suRLhtT2Qo50pkKM=",
-      "ops/lightning_attn/fused_recurrent.py": "SsJ+HC5opio7H9LsPWyrbnPOoK2EvOcxLuaq8Vp69GY=",
       "ops/linear_attn/__init__.py": "3OeyduLmyi1manyo7FQz1ZLSoLgEXAKbHKn1uqDcK90=",
       "ops/linear_attn/chunk.py": "5xqgOnPPA7nGPt/Rda4wOxSljFTD3Izw63VpuDLJNd8=",
       "ops/linear_attn/fused_chunk.py": "MD6FycgZSdn8jV88UDvxpioPAw6hCxvsSq49qNIKpmo=",
       "ops/linear_attn/fused_recurrent.py": "jUZ5gPuN2VngqbebHbN0lrIAOiZOYkxIku+C3Eh/9Vw=",
       "ops/linear_attn/naive.py": "46BX2exbemO1c8XXVhmkz2m6H1ywTmrvT4/odStEA8g=",
-      "ops/linear_attn/utils.py": "7X94oxSkHAGSpk7SQ5izs5Q0rYEDZZzUV/RorFdzY1o=",
       "ops/log_linear_attn/__init__.py": "qsAhfXSfoQBDmP6Li6+CjNd9ysMDGjcbYLonvcZcrjU=",
-      "ops/log_linear_attn/chunk.py": "NTaU1ErT6E0PzSt9JL9KayX8dQgLG++85PFwj3sRm1Q=",
-      "ops/log_linear_attn/naive.py": "xmXXuoIdEkTucVYJ/s4dTCvMWuJuO8KP0299J4jdtEU=",
       "ops/mesa_net/__init__.py": "32w1RDjxggXJ9JobAR01dDnN2UZSG2Fui7qePyQT2J4=",
-      "ops/mesa_net/chunk.py": "GTupTJ0BI3Ix2giVi7sMA2ocs71yvZ4uuKvf8rkUqKQ=",
-      "ops/mesa_net/chunk_cg_solver_bwd.py": "MvXMHU8MSOTNG62zURV+iMpvkwEh4eCo7epde0DIvBc=",
-      "ops/mesa_net/chunk_cg_solver_fwd.py": "KG4SXzD6dZomzRmYmqaWtmBoYVgSk6Fxda8uXjJc5go=",
-      "ops/mesa_net/chunk_h_fwd.py": "F3AUgFAYjqUAxslqpjsv+zW1PjdxGnwaU/Qjh5tT0iw=",
-      "ops/mesa_net/chunk_h_kk_intra_bwd.py": "17Le7w0rK0MrLUROFQqcZ9yCQfxpRYMhe5o3dRHGIfo=",
-      "ops/mesa_net/chunk_h_kv_intra_bwd.py": "mQnvHOYVu3HMX6IA4rD2zl+9w5BOi2yzFsr0GSo5+Nw=",
-      "ops/mesa_net/chunk_h_kv_intra_bwd_separate.py": "FOrwfCVxiMtLYaR9JIBqzxExwDWMWPVuOkaxgHD1AgY=",
-      "ops/mesa_net/decoding_one_step.py": "ZFuKB5H06F+uDagQLPGiu+Qk/qpZxjRbKcnuqWtC6/0=",
       "ops/mesa_net/naive.py": "zv3jT+71UegNb0JyIKnPv8OddyzHu4weEETtt0nJvzA=",
       "ops/moba/__init__.py": "P7ARNbLLMsmTtxcxPk09NPmjcvHGwWCmooXPYvT6ANk=",
-      "ops/moba/parallel.py": "8qUN2zo88Pldh809UyDowBVj8kCNplsS00fIxLRdSyQ=",
       "ops/nsa/__init__.py": "6Mw2zDl4Z74QXVjME1d88PMSKTSO/AYUa00lW/IK+eI=",
-      "ops/nsa/compression.py": "XjFqvEXvuIEeCBrJU+YzhWRH5UAkPGcUX4Tz/xx2M6Y=",
-      "ops/nsa/naive.py": "5A7lyoN39oK0Wc+Yzm/dQbNEtlaGmZfJsMDv8II+m7c=",
-      "ops/nsa/parallel.py": "BLaPXXsO8tcAJ79ltV3xwX/1qTganF2Nh+Lto/3zvyQ=",
       "ops/nsa/utils.py": "lbhybLC0CqXTn6liY8d1sQfTdZ2Gl+++Y5Go53OX4xE=",
       "ops/parallax/__init__.py": "g2pTOybF7cyoyaMAku1gpjV5MBY6s/Ll1X2NT6/w/YM=",
-      "ops/parallax/decode.py": "qHX0Le4hA66S2LLTcMa+p0VMB0uWvhYGI9urSPLrB4I=",
-      "ops/parallax/naive.py": "XEvAU0cSJtunA3ElSy4GFs6idarjK7Vb+xkrDKSp2OQ=",
-      "ops/parallax/parallel.py": "ZrLKoopSj60DSxmqbaS3MIf3AUyEpYh7evJPlTUN+AQ=",
       "ops/path_attn/__init__.py": "eyvdV5LkkxztzDTDYAl862vhXITUOcMtF6D8n6n0O/E=",
-      "ops/path_attn/cumprod_householder_bwd.py": "EZ0t+c2hG8Z/0IvrVxMBLDgO7zLgJsI9Bc7L0WYgvh4=",
-      "ops/path_attn/cumprod_householder_fwd.py": "UF1wvSIMRHeJmPf5qbFuxbC+qCm3ipAfeAVf3IEOA6A=",
-      "ops/path_attn/intra_chunk_preprocess_bwd.py": "HMQYfpKCHEMmzomrzb60u3gJK419yOAg2kx9XolFf6U=",
-      "ops/path_attn/intra_chunk_preprocess_bwd_prepare.py": "6XIZztbrFYFR/7XJSIy10RxKid3+6dcZ5c+A6XL2fBs=",
-      "ops/path_attn/intra_chunk_preprocess_fwd.py": "U8P46fLq//0vDN0SiE3cwI4n6FbAO6/UGNhfOiXegRo=",
       "ops/path_attn/naive.py": "4QsKiZUcF0q4IJbrku4Xu/IR3YBg+n6Epio7FAOW9lk=",
       "ops/path_attn/parallel.py": "V1c/T2VKbs8MgNJ0AT2JA5ULpx49aLb5maoxizAbVjk=",
-      "ops/path_attn/parallel_path_bwd_inter_dkv.py": "6/EOBq9NQdKaCrEGwPwbi69vKbOoNCsKU1q2vRv7YFc=",
-      "ops/path_attn/parallel_path_bwd_inter_dqh.py": "fp1NEXRSVAqhkmGcvoNfQn2pt5dtD0plx6W0TGR2xDw=",
-      "ops/path_attn/parallel_path_bwd_intra.py": "+3QE/gzK3fAYqgo/WNHNC5LV8NJrL4+UHjDN6wOd15w=",
-      "ops/path_attn/parallel_path_fwd.py": "cx6fGpfZmEuxHhl81gsLOBAYa61ZnkI9usqTakOIV7A=",
-      "ops/path_attn/prepare_k_cache.py": "LqEoXyDlVfIaJQhrmpoS3pmetIaavAm+SmIm1NKFohQ=",
-      "ops/path_attn/transform_q.py": "uZ9u6JDzXFKkWiim4KUkXWoAHYKlSDK5L1fMBNFrEGY=",
       "ops/rebased/__init__.py": "pawYkzOJ/mQY5Qj0Ibuq81P3KNgO9Lv6swOxCjlDlGs=",
       "ops/rebased/naive.py": "XOrMn2Ohd7MDJ7QvP93dPfjv/6prnZW/ozxHtU0KqWQ=",
-      "ops/rebased/parallel.py": "+3WRG3oxV0T78e8HWkN32Q8KBe0V2GMQ/OpxowWfXFI=",
       "ops/retention/__init__.py": "i6uHA1ndTR6T/fVfooceKo12wyBFGYLe1EAm0Mo/sg0=",
-      "ops/retention/chunk.py": "x0fX3U8KTjr9WLQK9+2C6Mn7cYNtACQ8RhFuUoCGpBw=",
-      "ops/retention/fused_chunk.py": "sAmIqFiH4IFM4yiM1G/pNvlsZxLtkORemYy5WW4i1pw=",
-      "ops/retention/fused_recurrent.py": "beES/zDbxCUWWWL9HSCc72a1+DWVvDKpBIsjKjkUet0=",
       "ops/retention/naive.py": "MG107SmOiRKLTZ8ziqaWN85nbgIPj/GqNJJ7KMt5Tag=",
-      "ops/retention/parallel.py": "gTnLcV5w17fZAOJyJttw52LGAe9XlkuS++lMGTZ3vRs=",
       "ops/rwkv4/__init__.py": "GnWakdtZiQJ0+fTYxqqIyosoPwjH6Yp8P1AbpHIDnts=",
-      "ops/rwkv4/fused_recurrent.py": "0fvq/4+NZMq6tzq7oHMUCaPg10qOzteSKLj5cmQKhvk=",
       "ops/rwkv6/__init__.py": "wPYgC6wG+xc4I2vessGhv1EaZ4tL5yuVQVSwi1rEEUk=",
-      "ops/rwkv6/chunk.py": "2VMAk6p0tKEudhieXLIg2bYg+89JbsGUiuhFpfvx6C8=",
       "ops/rwkv6/chunk_naive.py": "zyj0Og15jGxNMzmlD9FoEdf/odzu3tqio/jiceVMNk4=",
-      "ops/rwkv6/fused_recurrent.py": "wMTP/Y17ayzQXxXw+fN8P/FB5KGefLB4bB0wJ9kDg6Q=",
       "ops/rwkv6/recurrent_naive.py": "/lCgua1vSaJDqA4tOWWUaVxol5Axi3cOcjI+mJj7OAM=",
       "ops/rwkv7/__init__.py": "dPKBISLNnwlKb+xp/8ejsc2peC9gKQ3NedtMT5wQ1sE=",
-      "ops/rwkv7/channel_mixing.py": "+XrwuE2RjLxjsiKrdhssskBTvW/zL9RpyFp85uv5J5A=",
-      "ops/rwkv7/chunk.py": "9+rVTjACj0FP2Mnwo30DHWIHkpZG6YfYN4gE6891Oa0=",
-      "ops/rwkv7/fused_addcmul.py": "uUzh8wc75jBLJhCagQfeE/yor25Jw34d9kQQNsQWg04=",
-      "ops/rwkv7/fused_k_update.py": "CETn2jde9U8K8kfVihcyjARgpcDT9nN09qw9wM1QQac=",
-      "ops/rwkv7/fused_recurrent.py": "nnNrjgLS5nv9plvU/MuMjt4/vDcGD0Wax9Dfigvwmp8=",
-      "ops/rwkv7/gate_output_correction.py": "Ah6eqJrlTgFGaGI711+pMmWlxCOpvHrExx8H9U+nO0Q=",
-      "ops/simple_gla/__init__.py": "sSyIh8oSfhyZI+VnYUM/uhIp+HL+gDcQBReLn/nJ5OQ=",
-      "ops/simple_gla/chunk.py": "eA/L+T+42GiIhYu1oG0fM66QFGnzW9Uz8oc2SnLTKig=",
-      "ops/simple_gla/fused_chunk.py": "eCALWbi61vaG6zPewlPmATvjlidM9lso9xQC8H4xAlM=",
-      "ops/simple_gla/fused_recurrent.py": "Ab98sk1dnnFQmBmueO3D95vUs0zK8PC30ZRJYEGaIOM=",
       "ops/simple_gla/naive.py": "g3fAN++vd0aWvZcdOW7Y0boMYcs2mOld8kePWRbl2R0=",
-      "ops/simple_gla/parallel.py": "hB6Zq+allqAYg/H8/9kpk0l7bsZCK9IsITkfLIhqZQ8=",
       "ops/titans/__init__.py": "4t9OK1aYdW03I+sUlfUjj3ALTez1M/XFD/H6KzeZtGc=",
       "ops/titans/log_impl.py": "Y5tMWJgtVE+GX394DdxQHd6Yirrg3sLx6sinMdQvkAE=",
-      "ops/titans/naive.py": "NUmchQQpsxFwNaEpSQbwZ/638zSr0tuy2xGa3DD07+A=",
       "ops/ttt/__init__.py": "cFaLRxAWRsKjcmXh01bJbh5U59ZaFaM7DwYgm/KBAbo=",
-      "ops/ttt/chunk.py": "xk4Cw5LLMUKTauseXa+oybDIOp6A0e+McU/vqzp3KhM=",
-      "ops/ttt/fused_chunk.py": "KDdTpBvHj3jcIm2EkkNZTxWkhGbf0qvSQYEFkTwf+d4=",
-      "ops/ttt/naive.py": "IaJu/Zd1/PgFA5Bnbo9WQkpdz6BPZD0bl8iGKkWQlCo=",
-      "ops/utils/__init__.py": "eUiV7OAFYzGDyI/WtrSNltVrfU3crATGfKFWo9U87Hc=",
-      "ops/utils/backends/__init__.py": "Bfq8p7M9iRlUzfIfPTlvsKd/LMutWWtEzp2wfwOOK54=",
-      "ops/utils/backends/triton_ascend/__init__.py": "rpIPSLMvhpuf3aqxxDUGQrrq6dKKF5Nk2qsJl+H34Lc=",
-      "ops/utils/backends/triton_ascend/cumsum.py": "AOxEn57vvwiNwOaDcVK3KbP+n5IUGrsbUWkG+kkukBg=",
-      "ops/utils/cache.py": "pksJ/9P1Gu1UetclWd7pTLmur9Kl8CxsjmCm4CAUepw=",
       "ops/utils/constant.py": "uVK9zwOd7kM/DPSQk8cllx5nkYeK0K/KZVyCu0ZFdPg=",
-      "ops/utils/csr.py": "comXJRo/nULFkkZk556qHNxBfWeXk0EofLnQ69EIEjw=",
-      "ops/utils/cumsum.py": "6oQ1YDSIgMFZOZRn2oX8VKy/O+B5mbVWXl6146frCmY=",
-      "ops/utils/index.py": "BdattFUI0BkVExltE0Zxl/BNO9Ldp6hSUfU1zbM3710=",
-      "ops/utils/logcumsumexp.py": "7uEgeqUuLxi06jN/hOcTHvqJEITDGXECjNzk5SKPce0=",
-      "ops/utils/logsumexp.py": "7g4hdIjRKkO8CylqAs80XR02A5PyIWYxrKcc79qRcy4=",
-      "ops/utils/matmul.py": "TOLhNlPSm2JJeFmVsgh5D3t9LPUjw30w+mq5TIHef5Y=",
-      "ops/utils/op.py": "nFuDfGU88dXF1Pe8xrcxoda8QWZA5QTojiFsse/53fQ=",
-      "ops/utils/pack.py": "XT3/YxB71p6VO9oXPkpD0te42kbwU0dFaG3aZawdBHQ=",
-      "ops/utils/pooling.py": "LB2lncX1ThcjDIEWvMwiQpW2T8KE/AQ5dTqcJC/hNhs=",
-      "ops/utils/softmax.py": "uq/rygFZaAImpnHgOmAN9D+jrq/m3WHTyu6QGQR/EoU=",
-      "ops/utils/softplus.py": "tdlxnzZRHWEZWhVl9uFVl/EGQ6phOBMEEwlbCtu8jz0=",
-      "ops/utils/solve_tril.py": "oy0zat1/r4IjvqQ0MkJbrn1MYoayluYORr3nvX7z9NE=",
       "ops/wall_attn/__init__.py": "r/FbDGjGO8WK1DaMbl9axWHpENUa3pzSc0v9qviHJ6I=",
-      "ops/wall_attn/decode.py": "JphI1kWj0Wwu024gnwgUWrWmhkAnKloXi1ik6cDpRrY=",
       "ops/wall_attn/naive.py": "9YKwBSvQ9imy229T5KMYKpuVNb+xEzhXsXdMemRWqjQ=",
-      "ops/wall_attn/parallel.py": "oa2UJtSi8ouSG8g6bUWHwOehbwLH74SSHIU8LUiwVrg=",
-      "symbols.json": "0pgr+/ByS6ju3XG9D6VWYAcRncaie27x08XB4GIIXKQ=",
-      "utils/__init__.py": "bdtrr5b2e7dzcO88Il6bH72LypiM103tHkbY2kPoaeE=",
-      "utils/_compat.py": "kYAtqK4JdTczM1FpqT44T/mILN7DcnPb8C4uPnO4Ho0=",
-      "utils/_config.py": "l1PBnf2C6SABltOAhs9GB0GTr0olQ1fA7d/bhzTb6IA=",
-      "utils/_decorators.py": "8OxDii2b+D+deTCQLPK1sdFx7McS/Wi9+AO67mlLrYg=",
-      "utils/_device.py": "yXcM5ZkBLCOvHl/Dh25SveTgbSemR6PkFAl/UabLYWE=",
-      "utils/_testing.py": "eFchc/74QJDXWv5T1YmIZFzRRpFWh+3gaYpQ1doztzs=",
-      "utils/ascend_ub_manager.py": "TcjVL6L7YiNwsa7YZ4+9KrIPQ2RmG7PJCfmXYc24uEc=",
-      "utils/layers.py": "afB10gGysmmVfVRlPnduPpuvcdAeDAstmqPC0sYJbRw="
     }
   },
   "provenance": {
@@ -314,7 +355,7 @@
       "dirty": false
     },
     "kernel": {
-      "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed",
       "dirty": false
     }
   }

 {
   "name": "fla",
+  "id": "_fla_cuda_faa8674",
+  "version": 2,
   "kernels-minver": "0.14.0",
   "license": "MIT",
   "python-depends": [
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "eSkze5EKF+DhheKdXmGOOQPa6oJCFDY0fw3vg7WX3mU=",
+      "_ops.py": "q78tXrH/5M5hfjOclPcmC6NXw+I2de7vH6gfnXW225o=",
+      "backends.py": "h+9RflmubYjIEs0Fqhadj8SzFXHr+ghgXBmqn5Tm9v4=",
+      "layers.py": "MogF8nJKBsgqJjVXLjVFJHBMakhJCkK2of1ykPWF0yA=",
+      "modules/__init__.py": "rq5MxYJ+DgZLmrKdBb/PB1Q6PcuZ+rP0btubKE7uMdc=",
+      "modules/activations/__init__.py": "kjYo6ny71n6YPTd5Y6FrBPehfjS1Tl6rZh1QYcflZ9k=",
+      "modules/activations/ops.py": "toY9DY+ZmrAgJL4e8QOVKpYuw7QSYLUTxxCU3m+8ZaQ=",
+      "modules/causal_conv1d/__init__.py": "Z439ju3G3EQsurHFMVv0MQkKktuFw8/drnyp8oflj0Q=",
+      "modules/causal_conv1d/backends/__init__.py": "ZxTK7vZktZ7TBQB/5oiJr+oIypQU1cgkbn/ybwJl450=",
+      "modules/causal_conv1d/backends/cuda/__init__.py": "erS68kB9gBCaE0rX1ES+sKr8FCKE9pH/EB3rPWCrYCs=",
+      "modules/causal_conv1d/backends/cuda/cp.py": "s1x3wm+BuGbY0uAVGw+cLUBOtuixmZH80EFiyYLEIvQ=",
+      "modules/causal_conv1d/backends/cuda/ops.py": "73qKlHa75wF92ugakgiqbNYthaEDltGtAqz7ClVMsu4=",
+      "modules/causal_conv1d/backends/gluon.py": "KtMSmMOGSiioJGcWT237Uj27ryoZrsaNMx8BoOqAf9U=",
+      "modules/causal_conv1d/cp.py": "UGe4T/utfJ/mjhPM1WDnCJL8q/I8c4rpBQUDbRFjXyg=",
+      "modules/causal_conv1d/ops.py": "vEvfb5BXlVWOqXzzf5Rt0WcpEbjnfw/fYnOKxi91NmY=",
+      "modules/feature_map.py": "1+D3QXgIAEspjPM3aZ4ugVRSPeYEWJ+eqN3djUyQofk=",
+      "modules/fused_bitlinear.py": "cUFUweFZFmJS/Zl+8Ex3Kit9EHJJkcWabyH85wuwtDg=",
+      "modules/fused_cross_entropy/__init__.py": "nLRzNtHqTxjxgQT0EIs9RzmwfAu6baobksBU/Znd7rA=",
+      "modules/fused_cross_entropy/gluon.py": "qO/HNquDhMHIqN+59VlJcC37FwcqDaSoIz3mZMoH2Po=",
+      "modules/fused_cross_entropy/ops.py": "EwWAwIVHIOFrB1QPh335jjlvDYfVPEN1K9FjVtGxFWk=",
+      "modules/fused_kl_div/__init__.py": "L6dw1Iru9kTpqpwGWP81uMBUCf719A5eF+7awXroNYs=",
+      "modules/fused_kl_div/ops.py": "BDp6i7k/mDtwCi57T/ZKCUFS0lJpVm0vktiGH76bTy4=",
+      "modules/fused_linear_cross_entropy/__init__.py": "h0f6nBT3ezaBQExjxH+0ZpIfV/+jzp0cHNfvTkLQXYw=",
+      "modules/fused_linear_cross_entropy/ops.py": "qzLHd+ePGPNxynsHJiZXuQ5fGxTbhOR/q/n3ACAEZgs=",
+      "modules/grpo/__init__.py": "/+3ocSnRFVx59pEfhKGRhjw5SQctEOle3KwGSHUmhrk=",
+      "modules/grpo/ops.py": "58V3OfIwt4dBXMCvqaiBH6om10HKECwVeSHJ0h+XuzE=",
+      "modules/l2warp.py": "Tq8EvtlI+6mDqVhjd5tYYZgzCEuBCL0LJWWNFdKjjRw=",
+      "modules/long_conv.py": "EdjKO2tZyPuJVu7uvCXQTHINGySbcUb7lu6PK5kiKnE=",
       "modules/mlp.py": "q09Zw5Bttqq27hnlybbQ9ncr0jSw1UCrJH2Ewh9roV4=",
+      "modules/norm/__init__.py": "pjZst4x0d6Cupu9pETPPzEv/YIZOg2s5a2IOEwpAbZA=",
+      "modules/norm/fused_norm_gate.py": "Ex+3MGO0fgNz34LBJPxE77oL9JTT5Kw9TEUG6YUNAqs=",
+      "modules/norm/l2norm.py": "eSt9ICshJk489wkyGAPtmFqv678u6Iqj02HXtB+BthQ=",
+      "modules/norm/layernorm.py": "xqeuqNBIH79m7UcrlPlPjAZCzTc7vmfNsqNFBLH8xXs=",
+      "modules/norm/layernorm_gated.py": "6bzsO2HBwOU32pUuiiFwDsobrMNhABGPJguGCGk7VTg=",
+      "modules/norm/layernorm_quant.py": "BZHD/kPfNszPl0JsW9j+ZUOnHQOH2LwdaRNj9pv/o/A=",
+      "modules/parallel.py": "so6JnuagTE/nZBIjhNs/mwhcZtAbRr+DEwDCT3yLB30=",
+      "modules/rotary/__init__.py": "Vijv0n40cGYaaX/x3M1CsnLG6gjVxsmjTflZLxzNAq4=",
+      "modules/rotary/ops.py": "EGgH28oKJM5lcMXl/5tTD+j4Yxfq/L7n8s/jZA7rAUk=",
+      "modules/token_shift.py": "SiiXroq05nOI+fYGTpMxP/J7M3VnGonWr6eF6bWspI8=",
+      "modules/token_shift_cp.py": "q0OchsRLmUl6mSEdDgtJUpT5Y3HA4wrpUiSWHY3Hcyg=",
+      "ops/__init__.py": "c1AozqIP1VhM9A/lV0mO/nrHrTkbgaNn2aSyM6y0CCc=",
       "ops/abc/__init__.py": "WwLiIkYi3xJ5RiG/bOCNNRaeeqmMuIPT6n1g23qfpDg=",
+      "ops/abc/chunk.py": "UiMPJJgp/whOIdGxHx9ZkfKgbM8FCMEb9rd4oXXQ3cQ=",
       "ops/abc/naive.py": "AZ3uz9eDvycW3gTwZuv9J4k86KndfQMZiLBFyPmzsLw=",
+      "ops/atk/__init__.py": "7D6lvGGxM+O5oAsRP2Rs9x9fFh48dBE4RuanHM0lCsM=",
+      "ops/atk/backends/__init__.py": "k/wccBdiH8LsLHjs/LB5l6utDxQ1xwxehzK1wC1Cdl0=",
+      "ops/atk/chunk_atk_bwd.py": "KP9Ek4fmTbKxre3TlqhoC3EzkdnI0DbVw55/gK8cndY=",
+      "ops/atk/chunk_atk_fwd.py": "BRYqeWa2M0n2gvkNlFfIenhyr0hcMP3BcoeiojEqUQg=",
       "ops/attn/__init__.py": "gU/67mJonbeXpzfxgi8K7HIG2a8Xpjd0gOuCk/Pu9EI=",
+      "ops/attn/backends/__init__.py": "YSU1FRHkTNYuPlHp5Ky5ZC7OgD0Ul32qARLiXPf6qoo=",
+      "ops/attn/backends/tilelang/__init__.py": "VJdfMdLIoy9BWt49a6z0mnyGiFtIBmzHFdfHRjiBYrI=",
+      "ops/attn/backends/tilelang/parallel.py": "MvrzoFaTR/z+yljm5MA1Piq5CFIT36ku1bMVYchAQrU=",
+      "ops/attn/decoding.py": "J1up9AtGqemUJel+bHCbY5tXQBn/p/dy9K0nwlV1EMM=",
+      "ops/attn/naive.py": "Ns1vl3PBVEcW7+4kwW87ATCccIcghu24uWr5fDTS0s4=",
+      "ops/attn/parallel.py": "bNsWwVHS7LNWtcpZ6F2aWLglgCVKcIjIA5P4puUDYMc=",
       "ops/attnres/__init__.py": "/ljBk1VNFjDa9XABPXqIiQ05vUdzRV6Dbx/KkQihrWs=",
+      "ops/attnres/backends/__init__.py": "R08rRiTtJKcbPZBFsllcEpnbT0bSzQk5NGRgn779j5g=",
+      "ops/attnres/backends/gluon.py": "3TwxEZvH8NccliN4Cu4PzsYYMiqObQwIqYOaKmmdpiI=",
+      "ops/attnres/fused.py": "RxdxfbEKzeZ0rvieg3/4nbnV+011tb9ZtVEOT46mnfc=",
       "ops/attnres/naive.py": "nV4WPz/vfRwfRsIzF7pKV95fEx2dQP4BrrPoQCh0a0M=",
       "ops/based/__init__.py": "qdBLUigpgjQaIN3HTtq57uxoqAaBnQkZbFiEIf18Rvk=",
+      "ops/based/fused_chunk.py": "gBng4smV8X5eHn0pgo677bsu8kvOZ+DNEkJbvegBxZo=",
       "ops/based/naive.py": "YWIFSLeslp2ruyQDz/uCr5V4ao2cyRyeQH1wiVLMB8U=",
+      "ops/based/parallel.py": "daBATK+vtrog3XYFx2WkTfyg0QPH0gGb/6zlaTZCeIs=",
       "ops/comba/__init__.py": "FBKonbvPPFHY7an2Auh1sXVJLbcgwX8aAD5fxi2oecY=",
+      "ops/comba/chunk.py": "dtxGTXfgVY0c8TMq6uSBDUBjaP9tNk5WsnScEyVDVpI=",
+      "ops/comba/fused_recurrent.py": "sVNBWn5bq7PI7QIpCGbDMGBfl0CMzDzv767eR+5jCiw=",
       "ops/comba/naive.py": "QH967c5rxlSSDs6J6jRtLjZ3iLiwgAYgY9+0m9rgD18=",
+      "ops/comba/utils.py": "1nCCwG8f9KZlupZOwb1M3zH8MkUKILB44GaLFxI8cUw=",
+      "ops/comba/wy_fast.py": "9cijtR5+g55508HeHN0nedcRcBu5doCZad/6b+r+NsM=",
+      "ops/common/__init__.py": "yN9sU6JLmt6oHmOBursxys8xjykn/RubQlhzf+F4hfo=",
+      "ops/common/backends/__init__.py": "51Y0UR2Ro+8SU+rWoHNKpGNcBVC+F0CjQ4KjKMHA9cg=",
+      "ops/common/backends/intracard.py": "mExR+X9YFuNXkoM3x5t/78kdzUk/ypZhBq+0XWmEZfc=",
+      "ops/common/backends/tilelang/__init__.py": "N4GWBK2DAUoBSuxrTKD4b1Qsb1BtGHQZqBpXy57WhF4=",
+      "ops/common/backends/tilelang/chunk_bwd.py": "v6bPgT1lCK1MBdsarkfulW9rEQHEVYcZdyWh7k4F16k=",
+      "ops/common/chunk_delta_h.py": "bwziSs86L78fXsQNPMnXOSCVL1h/g80uWd3dvwtZuV0=",
+      "ops/common/chunk_h.py": "nR4KBApZUosIPbKg+rOZ/fL1C7vKJqowY6tUJy1pRu8=",
+      "ops/common/chunk_h_parallel.py": "vLG7BiBe4WhuB5M8x30rAeqFquGerzwIkRB2V5IXwVw=",
+      "ops/common/chunk_h_split.py": "bNGwtZ1QAA/qhJ6mpgDWMcIK/3043cTzyMMKUyFCNpo=",
+      "ops/common/chunk_o.py": "gmMl3/PC5JD+VjhlmnbcEivMAB9qsMvjhys6LFg+5r4=",
+      "ops/common/chunk_scaled_dot_kkt.py": "IHtve82Yy5FmP3N8Tq4FvCrlmoqzF5+4fjr7rRYi89M=",
+      "ops/common/fused_chunk.py": "R5PRqQBs8L+LDM3u15SXusRYnrJiqjYII168tWNbkDQ=",
+      "ops/common/fused_recurrent.py": "ZtFhTDcrB4kzHOEnqV4glzLCUf1x6RO14C+rYorz++Q=",
+      "ops/common/gate.py": "WOCYhSul4pkQvNUX7Lco3RDFsf4Dt6Qk6UN1f5Rc8WY=",
+      "ops/common/intracard_cp.py": "FUJFbWiaKt/APVmS0TQyMrHZEVAaAi6eu/VM4jhBqRk=",
       "ops/cp/__init__.py": "uud2WkodJvH5U5DjBa5oSLJtZ1G8O4D4XEuspXyC9lA=",
+      "ops/cp/chunk_delta_h.py": "OHQD+CubC1RfbNqxsRewZVh/NSNqUMEOuplaNm9FI7U=",
       "ops/cp/comm.py": "lqyQDwmK19vowTtGJGO+6SEn/neFFE0G8xFvlxLOmb0=",
+      "ops/cp/context.py": "VS/fHKQ3dAhML1NkHpjch48iC1KJjE2KTBkL4rVlbF4=",
+      "ops/cyfa/__init__.py": "OAx8kAAAHpFVQxe3+SV2psaScEIJ5kbnAcOJkOrWo0s=",
+      "ops/cyfa/chunk.py": "QgjD/Bv7TlL0qLd17SLVWFs4xas+w61XQ+Zfq1yoPJ8=",
+      "ops/cyfa/fused_recurrent.py": "XWLzv2l5TOjGiEcx46f94VVcErCOOMO0+vDO+dn7GdA=",
+      "ops/cyfa/naive.py": "95GVYqF/wXwBfLsMf9x9lwJLiSk9LfPWHQyl1QGOYn8=",
+      "ops/cyfa/qk_norm.py": "phhdmmvyB6qRosmZC/ybLBKsb6UjXQj/QFU7gjoYe04=",
+      "ops/cyfa/utils.py": "U+R2Biu1tPo5GoIEngh2T1GebIgI1yK6gEDfX8qqNGw=",
       "ops/delta_rule/__init__.py": "1gIkBajhtbVwuJOgPQ+FYYdq4tYgNgib+veOnont+N4=",
+      "ops/delta_rule/chunk.py": "ueCVdrAgOpCxVVb7gDgihDRLQck9dY0dOQB/qdxHxGc=",
       "ops/delta_rule/fused_chunk.py": "2NiUBRAl98YgRVGr2uzYJV22pQW3WNOe6euTX1Mdqyw=",
+      "ops/delta_rule/fused_recurrent.py": "WEg6WYVHkdvmncKPtXfYky4skqIe5whzjrX9rE3PAHo=",
       "ops/delta_rule/naive.py": "IK4DL04IAZtCmMyfzyw5fYcfqsD9HxgCQ9n2makZKWg=",
+      "ops/delta_rule/parallel.py": "wR9GWtXrKRwFP+i1ewrvSuO9OJRq+HwYeAkF2pwlchM=",
+      "ops/delta_rule/wy_fast.py": "voSaPCMCS/G8crBpowSqB1KI1lQZZQdNuyRlmDTiGVI=",
       "ops/deltaformer/__init__.py": "51ZmNFoHy3P0lVn/2EHKu1QS2BMC6h+hnAwSIr9szYg=",
       "ops/deltaformer/invcum.py": "ren181AxlHBv2MoZIHmSGMmJ0sW6HDb3t8fqG7gDELU=",
+      "ops/deltaformer/naive.py": "W78CzYUyuxLY9fEkKghSh2dnjux0NOJ6pAVhIN7M4PU=",
+      "ops/deltaformer/parallel.py": "XMKE9sLFJl4UVspOXQxhefztzykDEUp7+S9GyUhmOwo=",
       "ops/dsa/__init__.py": "ZNyws889LPBiY8gwB5zTs2vQwJYV8bX5jg6Cldkkq9s=",
+      "ops/dsa/naive.py": "JKyMUzwK44XrwlKUWfLhU/iK5DlCsXceJuxtXrzaCQ8=",
       "ops/forgetting_attn/__init__.py": "A6ozg9zdEsz4lgUAXqEIUY8CeNp+5IZqYKOpChpj/A4=",
+      "ops/forgetting_attn/naive.py": "rjmhL74ja9e060lOrcRgMpAuwO5xDqBKbXw+YoE26+8=",
+      "ops/forgetting_attn/parallel.py": "nm4vnTE92QQIvgF+q80TGjweTw8YFX0tyN4Ip4vE7ls=",
       "ops/gated_delta_product/__init__.py": "0MAzYdxP0ODZjN4WNQ+3HLwigokyfMv15rXUpTX+qnw=",
+      "ops/gated_delta_product/chunk.py": "jbQM/fXoRX/nzNZdSgUB5NBB7DYzhQcabsDpuMUUbHE=",
+      "ops/gated_delta_product/chunk_deltaproduct_h.py": "w/+GVM+h+ffj5shd1n/rPqhhwIfccadOiUQPU6slGYs=",
+      "ops/gated_delta_product/chunk_deltaproduct_o.py": "HAmJ5oygn3R7gHe7Ur5kWAFf9sV3qtS57ZjML69FU34=",
       "ops/gated_delta_product/chunk_ref.py": "R1dHdPy4rU+lTQQKodAn9h2xPLxmwI2z0T6pRvh8Llo=",
+      "ops/gated_delta_product/naive.py": "bJ1Cc3mljHcAXjc4w1JDyxj+i4hmhCPsoELKrTkyvIM=",
       "ops/gated_delta_rule/__init__.py": "Xt0hhHr3CECyZDmBIZakdufk3bvvIa/V/iBmpkYwpnw=",
+      "ops/gated_delta_rule/backends/__init__.py": "iwpwhvQepjkDHC+SDA2b3SziHFeegB79Oz+pU+TsTl8=",
+      "ops/gated_delta_rule/backends/flash_qla.py": "b14UI75g/TZIsJLMQk5j71B4B2+JkmYL24zACdtDcRg=",
+      "ops/gated_delta_rule/chunk.py": "FJbbVRckUs3kIRXGbCppJdyOCP6/ThpF16ANRmSuiTk=",
+      "ops/gated_delta_rule/chunk_fwd.py": "35ZNk4cZW3CWG33fryWIcXLtJ4WGvQvKNonONIKCvWM=",
+      "ops/gated_delta_rule/fused_recurrent.py": "b4i7Wb0iuRjzwlXqmAb6AAZiW2o4hWVv5iDJ5dAuGpY=",
+      "ops/gated_delta_rule/gate.py": "qAq4nKirRZtlYtHYYbi/jXC8HZZiRaftCwu0haEgmhg=",
       "ops/gated_delta_rule/naive.py": "0c8XmSNJ/T6Ur5mbIuPTqBvkotGIHOW3CjRXJXFm4Ms=",
+      "ops/gated_delta_rule/wy_fast.py": "IayqzhuCGiM6lbjjWKAi1GU4IvgXnZr++xSNq5ZfSKY=",
       "ops/gated_oja_rule/__init__.py": "oShC/cL0jMy+xWR7YBV9QxkAdxt1xasXDCypGAU6Hgo=",
       "ops/gated_oja_rule/chunk.py": "lq+5BXEaJLfL4dqqcJb4TmtyF9OrP9Wh+JpRfpdzd2s=",
+      "ops/gated_oja_rule/chunk_h.py": "K6g7OMKN9/uUqboAelMpsu8jCwJ2opTKebyZiAEfIYE=",
+      "ops/gated_oja_rule/chunk_kkt.py": "9K2pSMxQ3FkmGgiH3JqbN90nixJymRK/x59ohwZYOKk=",
+      "ops/gated_oja_rule/chunk_o.py": "AxQIEqnIrWZQmuniM8gk1h+Ynd6d5saoXeJigp03ORw=",
+      "ops/gated_oja_rule/fused_recurrent.py": "iHdQaLJGJSLM8X6fYmOqtVe07o97YvidynXGewNG008=",
+      "ops/gated_oja_rule/wy_fast.py": "/+qHRfD2tvWfOfPHa7OMdFWezHLqjFnGJyhiUbF02u4=",
       "ops/gdn2/__init__.py": "xiEsgvVEfHhKcurQqL7ZKGGzEBW2BtGAylE9vxoNmAk=",
+      "ops/gdn2/backends/__init__.py": "d3Zf1qLfl/RHs8eAs1cGjN05HBvbxvoWdMke7s1BatQ=",
+      "ops/gdn2/chunk.py": "y4dx8rK1CBrPDQthDgnee9WAdWr3JYL5t+9W1/XKAbA=",
+      "ops/gdn2/chunk_bwd.py": "CGaUEbBKHR+yZJH8ILgOGDhbEg8S6i4SaI9hV/uXFow=",
+      "ops/gdn2/chunk_fwd.py": "GghmT44Q/8kQ1Jc10sGYeAcaKxmhldjJex2Gr/FSyec=",
+      "ops/gdn2/chunk_intra.py": "8wEbCvVZ/OTPPDyo9cu7R0WeRJZrqZqqLOAAFL0c2do=",
+      "ops/gdn2/chunk_intra_token_parallel.py": "W+tiIjr7D+nxq/4s1itDolgS7mya+LC3b9iJ1oDpMd8=",
+      "ops/gdn2/fused_recurrent.py": "TDAjD9IqdASKQX9FLLfSLFgsdTbDoE4jwnShmBQSpAQ=",
       "ops/gdn2/naive.py": "amWAVAHrMA34W2KXYRoN0VjdGDQjzyvNXg8Ozd4+nno=",
+      "ops/gdn2/wy_fast.py": "N0XZaivjPUDGmWPS8tyR/2tqedUcMtlHpaZRZQgn72M=",
       "ops/generalized_delta_rule/__init__.py": "cn6rKWEDrxWw428Dk6EEy4Y7nyWHpc8rtR+EbRO9ndA=",
       "ops/generalized_delta_rule/dplr/__init__.py": "F8TbaiZeRKjbKhlI4XYopNd4bfzV10feNO06i32h2Vo=",
+      "ops/generalized_delta_rule/dplr/backends/__init__.py": "VkuoS48+KY513u/s51Qpx/Qp4oSuzcMgoeqzc+9uqWE=",
+      "ops/generalized_delta_rule/dplr/chunk.py": "KF/hPswvDmMtN0+Q1OQR9TbpDOrMqAWSs5qPbuBOscw=",
+      "ops/generalized_delta_rule/dplr/chunk_A_bwd.py": "Sf4K4R7jzOqgTCo/fLhoT5evBHj03sn1qKuqYZLpt3s=",
+      "ops/generalized_delta_rule/dplr/chunk_A_fwd.py": "szc1lBh7LvMrdXft3OQWQ+vklQxDrqjmet3d+69VAr0=",
+      "ops/generalized_delta_rule/dplr/chunk_h_bwd.py": "/rB+pwDapn9SM8xqp1x3Ro+N4EuvgIoYhte9f1Wl0Jc=",
+      "ops/generalized_delta_rule/dplr/chunk_h_fwd.py": "KXBpSZVK4Qvp/2XTe8TVXQw6M8t06N3AHdlldP1+Yi4=",
+      "ops/generalized_delta_rule/dplr/chunk_o_bwd.py": "9NLZBaPZ8nlLSfFbbVpoUz3fM5H/VR9nRajBlY2Nx04=",
+      "ops/generalized_delta_rule/dplr/chunk_o_fwd.py": "2155nmxibRXn+jOz2MGsUnkdNhWsaHKFvOjmBAmcIq4=",
+      "ops/generalized_delta_rule/dplr/fused_recurrent.py": "QT9BmtI4OVj/Nh6Qo6UM3TYb9W+ZifSWHVnJFGKI4Uc=",
       "ops/generalized_delta_rule/dplr/naive.py": "4flLSZUtwaq5S/PH+fu04rzOM2+USdRX6Ofte4lanE8=",
+      "ops/generalized_delta_rule/dplr/wy_fast_bwd.py": "fcpJQxkNls5tEKkP9fxaSqrGkrdZyNzmruFfGpe7hCY=",
+      "ops/generalized_delta_rule/dplr/wy_fast_fwd.py": "aNodADVGx656FWS5f9e02Lb0lrp4syVsW8NR6zYZF1Y=",
       "ops/generalized_delta_rule/iplr/__init__.py": "6Vg8MTc+Byrnmx3HJz2yySi9nk2kyKCSg7aL4CMJ+GA=",
+      "ops/generalized_delta_rule/iplr/chunk.py": "YMAVurj1Jmp97+sWuTlzfUOaCUeKlhQzbSGXTzobwW4=",
+      "ops/generalized_delta_rule/iplr/fused_recurrent.py": "u5LKf/1u7v9MSnb3mixwWa6/5h7GdDBJ5CHf35Tv1gQ=",
       "ops/generalized_delta_rule/iplr/naive.py": "YotccLOueXwBcQys4ICmFAysSrX8fj9eRasXIpmpRMo=",
+      "ops/generalized_delta_rule/iplr/wy_fast.py": "hqVfLtmdlhQnlIeU7LrByAxBcNYtVSehVFPZW/o8ZhY=",
       "ops/gla/__init__.py": "PTUqqQlC4cXrhRhG/KRx+7X3m9WJKxGyzrIB/WJEXSo=",
+      "ops/gla/backends/__init__.py": "3LkZ6RPx/+FnCqaj6NQC8FovyO0XKlLVkrKopGEt+nk=",
+      "ops/gla/chunk.py": "5xTGVCzeQzxbSf8gg3HdcpBPCqOfMykToxXU1NI/S1g=",
       "ops/gla/fused_chunk.py": "OjzKcjSQujVQ+Q+d1Lmg4tsTzZaK6xWs8AGuplRn5zo=",
+      "ops/gla/fused_recurrent.py": "ToICyjG/Jl9CPbYKh/8XCmd+tC0Wuk+zQRrDczQ1yCE=",
       "ops/gla/naive.py": "rEuH5bEAieya1CTvtZbebVFo2J5LdOX1ztL1WiDkZB0=",
       "ops/gsa/__init__.py": "9uygbV6orNtaTkZl9Hz+aAp7EvT6RuffawfeFtuUNEE=",
+      "ops/gsa/chunk.py": "lAt7ynwdo/OuOmueHQR05KhPgCO+RtZzpPgQf8pz73Y=",
+      "ops/gsa/chunk_bwd.py": "zph8XOxa2WmwkAIUpIuLTjpSubQSp8uNiphJeeW+kfM=",
+      "ops/gsa/chunk_fwd.py": "goDc5m91bXe1spiElAX1fM5R04qTQgbbvdM7GSc+rsU=",
+      "ops/gsa/fused_recurrent.py": "BOxAXEEw6TwbMVpVLR98bB590Y6CrEOfP7/xFu1Ov3Q=",
+      "ops/gsa/gate.py": "+oIWuMo420dfL5QIwk21VCjkohKKmwtmoIEVeBdrk9A=",
+      "ops/gsa/naive.py": "8ECW81ylA8gLsR3svJsPYd6pAzW8ELmSMuwZzkCJmeM=",
       "ops/hgrn/__init__.py": "69DMv7FQYvobQTXqOFtu7muAt1k1k9nvsFXcMAmjQmg=",
+      "ops/hgrn/chunk.py": "odEz3XdQ1whBWOhYQfCVvMS+zy6QZ36hb9USuqjwV6w=",
+      "ops/hgrn/fused_recurrent.py": "SddIcmnP1KUcGWTk2rJGtPEoM5CnuV+zzMOUmK4pc98=",
       "ops/hgrn/naive.py": "XKo+5zJ+G4elkgK5ARsoXm5WYkG1QidTKNIuPT4k588=",
+      "ops/kda/__init__.py": "OsCs1naQ1tZfkuT3f/9QXuufPv4RKCM+Gya7U8easqU=",
+      "ops/kda/backends/__init__.py": "TtEq1fz8FfYrdqbSdZHPA9KNZXCgERZ8sINxDCYZxN8=",
+      "ops/kda/backends/flash_kda.py": "ncJlqPB9icS0hxx1gLWACHX1u9gWK9Zuz2OnLW7qSp8=",
+      "ops/kda/chunk.py": "Tc5Wt6dCdvv3AyIjn4HyiSNQ+PEhlz4vXd9oQOyDCro=",
+      "ops/kda/chunk_bwd.py": "/LwpwY5YXbqM7QLO2t0PogkAptPrR0y/tvRqVd0E20o=",
+      "ops/kda/chunk_fwd.py": "yPLK8XehVLVWe2dl5sX9eFjjxN5G1YX1rSKRyu8zLhI=",
+      "ops/kda/chunk_intra.py": "0iplpMZqGD//G1sLbMdY/hoksjAJu3qv0aBoOfHrNio=",
+      "ops/kda/chunk_intra_token_parallel.py": "a3dvImIfizoopyQ2i+fKtetL7001HavfOQENgqw16/I=",
+      "ops/kda/fused_recurrent.py": "KWsqM6LZXU2n7faYHMYMqwR/N1byanrjWc0R3Eogt6M=",
+      "ops/kda/gate.py": "sHa64RkSgT8vMYUiG0fYxmYQJMvt4mG2TKYCm+pa+nI=",
       "ops/kda/naive.py": "YKMihdS2cGj/YztIu+irMQKAZtJPANJ+EhmaiPxz8BY=",
+      "ops/kda/wy_fast.py": "PasjSQorAguMGZkPwmVtosiTUlZ8x9jLiCln/UtX7bw=",
+      "ops/lightnet/__init__.py": "xI7MhtSqpNXoMRr9nmXTmq8swM+XjxBEsMH+9VKqwqA=",
+      "ops/lightnet/chunk.py": "7He3VzwLRiJ7avtCMje5Oz+nuRkMmtMEzdFUxLyEh04=",
+      "ops/lightnet/fused_recurrent.py": "hzn/dfLsBVrfpO44OQhreQIzysWT0LYHBoIubvphmVc=",
       "ops/lightning_attn/__init__.py": "FPVVYNLjfO3GJQYIp+4HCQ4MdJtF8UZx5yktWDIatFA=",
+      "ops/lightning_attn/chunk.py": "w43Qdzp+ogeMkGAOL8ZvxVIFV/nBtlcN+5vn/oBhZCA=",
+      "ops/lightning_attn/fused_recurrent.py": "w98eqzLvLE4Gs8TR4iZvMl7X8mGjmL2+jSqTu0h3B6A=",
       "ops/linear_attn/__init__.py": "3OeyduLmyi1manyo7FQz1ZLSoLgEXAKbHKn1uqDcK90=",
       "ops/linear_attn/chunk.py": "5xqgOnPPA7nGPt/Rda4wOxSljFTD3Izw63VpuDLJNd8=",
       "ops/linear_attn/fused_chunk.py": "MD6FycgZSdn8jV88UDvxpioPAw6hCxvsSq49qNIKpmo=",
       "ops/linear_attn/fused_recurrent.py": "jUZ5gPuN2VngqbebHbN0lrIAOiZOYkxIku+C3Eh/9Vw=",
       "ops/linear_attn/naive.py": "46BX2exbemO1c8XXVhmkz2m6H1ywTmrvT4/odStEA8g=",
+      "ops/linear_attn/utils.py": "tjqU7ST991YBV68Zs83Fy/li6sUSMeKNPGE3VcYMvfo=",
       "ops/log_linear_attn/__init__.py": "qsAhfXSfoQBDmP6Li6+CjNd9ysMDGjcbYLonvcZcrjU=",
+      "ops/log_linear_attn/chunk.py": "21c+pLRtfBGC6RPf30KHsqafmCMBOcvL1SbjvNVUT9w=",
+      "ops/log_linear_attn/naive.py": "Ju7vJ7L6j1xSt2OrDHb1U8lXQ/G5c8TplDVwv2OGH6U=",
       "ops/mesa_net/__init__.py": "32w1RDjxggXJ9JobAR01dDnN2UZSG2Fui7qePyQT2J4=",
+      "ops/mesa_net/chunk.py": "53VxbqZux86q96Jk7ZT7iHkLBlcADKp9yZuvBDovRhA=",
+      "ops/mesa_net/chunk_cg_solver_bwd.py": "6dUL+GS455oU+M2QOvROmJTlkGIyiNV2qD2Z/KmlY8I=",
+      "ops/mesa_net/chunk_cg_solver_fwd.py": "2xUtPW36YmUfqpj1qQblJpyYlzLO/LXnPWbH4cHyLoE=",
+      "ops/mesa_net/chunk_h_fwd.py": "kUhk8XVi9GPrfuViwL4shPieN3ZCZs9hvqaT7L7Bf7g=",
+      "ops/mesa_net/chunk_h_kk_intra_bwd.py": "TzV3RiE5k3/C38J30lgLcoHHj1fWiT+39iHlZPrNLvM=",
+      "ops/mesa_net/chunk_h_kv_intra_bwd.py": "b4s2qkF5ANeggNE8HrNZT8LEfd02FqTLo1+KY4wJKos=",
+      "ops/mesa_net/chunk_h_kv_intra_bwd_separate.py": "6g3MpZX8313sLrlMhGuwA7MY2e+wVX56xgwqoWj3OZA=",
+      "ops/mesa_net/decoding_one_step.py": "D3QQNRDvQkfFnQ5CV77pK32GqM4nI8h5QqBgH95uBdI=",
       "ops/mesa_net/naive.py": "zv3jT+71UegNb0JyIKnPv8OddyzHu4weEETtt0nJvzA=",
       "ops/moba/__init__.py": "P7ARNbLLMsmTtxcxPk09NPmjcvHGwWCmooXPYvT6ANk=",
+      "ops/moba/parallel.py": "vlJbZdAJkBvuqbgjTD03lQvB5BK4kagpC1PCMYQdO0c=",
+      "ops/momentum_delta_rule/__init__.py": "5OaHW/StrtV3AXsAel5gY1tFYo/vIktNiATrRNV6C3Q=",
+      "ops/momentum_delta_rule/chunk.py": "PMijn0zk+fco7T3jTPV3AiGJjeL3c9wyxYGozOFRU4E=",
+      "ops/momentum_delta_rule/fused_recurrent.py": "gUZQM29isdhTDKa4eKL/lvzr2fZwe/iu2saM6hvanJo=",
+      "ops/momentum_delta_rule/naive.py": "p6i6fC2NbwyGyAa7BroNb//NKET5cWPuciubVEYq2Ws=",
       "ops/nsa/__init__.py": "6Mw2zDl4Z74QXVjME1d88PMSKTSO/AYUa00lW/IK+eI=",
+      "ops/nsa/compression.py": "XxTFQ7YicVLJ/RajeCgkQKUPKDq+t03lMtjgirTk7Mc=",
+      "ops/nsa/naive.py": "zJ1lQnQEvnF6G6j/n1HF31DDYPQOENL/fJfBlWM4a+k=",
+      "ops/nsa/parallel.py": "SQEcy8p6KEX+MbIRY+/o9IblXJPRTcMOMdJk9i4kSdo=",
       "ops/nsa/utils.py": "lbhybLC0CqXTn6liY8d1sQfTdZ2Gl+++Y5Go53OX4xE=",
       "ops/parallax/__init__.py": "g2pTOybF7cyoyaMAku1gpjV5MBY6s/Ll1X2NT6/w/YM=",
+      "ops/parallax/decode.py": "RMgM4i3NxaiktyPLxuF2wgPytE6Os7vbywv9DovgQVM=",
+      "ops/parallax/naive.py": "Bul9UJn1oJgW3Q2R04b/aQ87GIBx8Pu+DMalkXdmf8g=",
+      "ops/parallax/parallel.py": "U1/dvoFO4sXS3c13BHbMzbdUdM8rACR0aoGuS7cXxGw=",
       "ops/path_attn/__init__.py": "eyvdV5LkkxztzDTDYAl862vhXITUOcMtF6D8n6n0O/E=",
+      "ops/path_attn/cumprod_householder_bwd.py": "aheXzKc/5z9ug392HO9RvnRLfuoB2IqMDqJGCoQHdW0=",
+      "ops/path_attn/cumprod_householder_fwd.py": "Z+994gihtCIHIVOQDpsZG40KAY2h3QoPhRWSZ5xTM88=",
+      "ops/path_attn/intra_chunk_preprocess_bwd.py": "vznrWNVjKSVqJnroNX0rdTzZubF6gzBfIo30UKH/EAs=",
+      "ops/path_attn/intra_chunk_preprocess_bwd_prepare.py": "m8JxIsdrK/0RQgoA4pCpdvjTEz61TI1yZ2ga7H6NZfE=",
+      "ops/path_attn/intra_chunk_preprocess_fwd.py": "Dl+JNpXxX2WYfFWC6pGWSDnN1yJPJdHVLoJ41t/tgN4=",
       "ops/path_attn/naive.py": "4QsKiZUcF0q4IJbrku4Xu/IR3YBg+n6Epio7FAOW9lk=",
       "ops/path_attn/parallel.py": "V1c/T2VKbs8MgNJ0AT2JA5ULpx49aLb5maoxizAbVjk=",
+      "ops/path_attn/parallel_path_bwd_inter_dkv.py": "mtVBfPgeTWxQWjcQfee4Kk/KRmck4kEwADFuT+PfcOg=",
+      "ops/path_attn/parallel_path_bwd_inter_dqh.py": "moIY1PhioPDmbWQc7T37Tx25kKn/1OrDGPsBOai1xTg=",
+      "ops/path_attn/parallel_path_bwd_intra.py": "w1sn58vCEin95acLyFmq+XQDQbGiiIU6vM8J/Xp5faw=",
+      "ops/path_attn/parallel_path_fwd.py": "uRUFU10eJB9ZCDcVn4AYM5+JdTRdKPKh7z5cpEEE+Ng=",
+      "ops/path_attn/prepare_k_cache.py": "DOuhtk1NY0PG0WSTxQcc8EvY3zFvTFRRLkyd9pztOaI=",
+      "ops/path_attn/transform_q.py": "qpu6Y3/BBj9K4dYP7Nj9QB/9zVOc0FHZY0NumyoyZLo=",
+      "ops/precond_gated_delta_rule/__init__.py": "E05uI9A1ocZvCySI4+nu41mCPr4MJmZA1upU7ruAHM8=",
+      "ops/precond_gated_delta_rule/backends/__init__.py": "hoBAVLZQ+jOiVYJWE5G4A1doL5gqeqP/Dz0istkBoeg=",
+      "ops/precond_gated_delta_rule/chunk.py": "d6aY4870cAoSyBkY063zvl57SHM9G31OZ6OUYL80OaE=",
+      "ops/precond_gated_delta_rule/chunk_precond_kkt_fwd.py": "4rEy83SMB79qql+FGOl5aQ66OYTj6D5vtNEgjcJEBrA=",
+      "ops/precond_gated_delta_rule/chunk_precond_wy_bwd.py": "eNSvTIyY+LfG5O3PLIasH3+vEsVg9uFYeo5f2Wjlwmw=",
+      "ops/precond_gated_delta_rule/fused_recurrent.py": "xRZ3yREutRIZ2qFVeS5u4QCy3byrePYRhL859UenVG0=",
+      "ops/precond_gated_delta_rule/naive.py": "O6pnpfNdxyMGmOPxdh7IZ1ExMYwV1KJ+1/L84R6Eteg=",
+      "ops/precond_kda/__init__.py": "KA5gjXPBlmGs6vc1Srg4zngR+5FXv+O2a0VJJDHrY0M=",
+      "ops/precond_kda/chunk.py": "KUpUtIX0ogoKnxatK5SMPdJ7Ms3lX5zvbSd4Y70Jm10=",
+      "ops/precond_kda/chunk_bwd.py": "pwNynFCythreiWjz6ezuDjZ2eqNdTmYj1PqysBdWCSc=",
+      "ops/precond_kda/chunk_intra.py": "gcm/6QihmSSHoN4DhSqNGuY2RPrbKp60ENBODhsi30o=",
+      "ops/precond_kda/chunk_intra_token_parallel.py": "ayBWo/wqF/AiU9qLC7jHmZbk04iGBcd9/FNeY7CuHDQ=",
+      "ops/precond_kda/fused_recurrent.py": "+RTPNLwCt8JVTL54hHmRqvEHLLJeLYn+a/9VKrzD3vY=",
+      "ops/precond_kda/naive.py": "7Bj6uTWYmCxBRHuOCzs+PhkWtxlFsTGkVXAzZGzhh8U=",
+      "ops/precond_kda/wy_fast.py": "pj8DLVEG5/XHWFj9SICxTKphtgmrz4Zvu+giMq8NFaY=",
       "ops/rebased/__init__.py": "pawYkzOJ/mQY5Qj0Ibuq81P3KNgO9Lv6swOxCjlDlGs=",
       "ops/rebased/naive.py": "XOrMn2Ohd7MDJ7QvP93dPfjv/6prnZW/ozxHtU0KqWQ=",
+      "ops/rebased/parallel.py": "cBnn/KTq9FlWJIpc6Jj9uJ7J14VfGFadrxABo4bsoQc=",
       "ops/retention/__init__.py": "i6uHA1ndTR6T/fVfooceKo12wyBFGYLe1EAm0Mo/sg0=",
+      "ops/retention/chunk.py": "LyunHUb/Y0ItUN9EDTaix4N3Zd2e/7D9AoOQTSXyAMM=",
+      "ops/retention/fused_chunk.py": "Nzn6nYHcMyjL8uOWkwD10Htbn28xhomLMG4XZ0p7ru0=",
+      "ops/retention/fused_recurrent.py": "WGV6AAdpoGG79XgHWkjDkFWY9ZpvE99xmqwre64KuZA=",
       "ops/retention/naive.py": "MG107SmOiRKLTZ8ziqaWN85nbgIPj/GqNJJ7KMt5Tag=",
+      "ops/retention/parallel.py": "WRBtqBxCWpIZ/Lm2u3bXwnbeLbwwQ9ZrrajeYeNBoTg=",
       "ops/rwkv4/__init__.py": "GnWakdtZiQJ0+fTYxqqIyosoPwjH6Yp8P1AbpHIDnts=",
+      "ops/rwkv4/fused_recurrent.py": "LG9PdTsubmJvgqvbDvT8N+9y5SWQZZXfTxC0Hg1jU5s=",
       "ops/rwkv6/__init__.py": "wPYgC6wG+xc4I2vessGhv1EaZ4tL5yuVQVSwi1rEEUk=",
+      "ops/rwkv6/backends/__init__.py": "3Hgm1GsMIil45jzcS+4PB9iZi77lM/W+P4hBlV1+fj8=",
+      "ops/rwkv6/chunk.py": "scQIA9mDVrEFjX5IoVNG7WRDmr9ss+cZBTqLDVkiEKk=",
       "ops/rwkv6/chunk_naive.py": "zyj0Og15jGxNMzmlD9FoEdf/odzu3tqio/jiceVMNk4=",
+      "ops/rwkv6/fused_recurrent.py": "Nriss1l8NCBYj3y/tIJJLEScxNo5BUXQcqOYe2T4Hac=",
       "ops/rwkv6/recurrent_naive.py": "/lCgua1vSaJDqA4tOWWUaVxol5Axi3cOcjI+mJj7OAM=",
       "ops/rwkv7/__init__.py": "dPKBISLNnwlKb+xp/8ejsc2peC9gKQ3NedtMT5wQ1sE=",
+      "ops/rwkv7/channel_mixing.py": "xJAzrxM8dRByDCC9DdLRi2Q56AzxDKWPcpg5nhyo4wU=",
+      "ops/rwkv7/chunk.py": "jiQCm20KhWEfM3MT99MCmm3LJZyWQSetiASzBEg5P4U=",
+      "ops/rwkv7/fused_addcmul.py": "ZQ437Ks75Hw3qa7F+iV+dx76TKl7nGcyJd0IEvkAEIk=",
+      "ops/rwkv7/fused_k_update.py": "NuR4O/dI+4UHVqwKLxgBBsXC7SCaCpTAed6cY2cy4qE=",
+      "ops/rwkv7/fused_recurrent.py": "EUJaU6go5WjhYB9WrAhOriB04cr5nuvCJwGuxNDp/1U=",
+      "ops/rwkv7/gate_output_correction.py": "rwkLmmd3vqW8dvoJIR7f7uz4iR8PgEWBMnQbLb582k4=",
+      "ops/simple_gla/__init__.py": "qHTqgCNXrTUUrTSo+ZG6B553XTCDonl3K58OpQ7YcMg=",
+      "ops/simple_gla/backends/__init__.py": "ZxTK7vZktZ7TBQB/5oiJr+oIypQU1cgkbn/ybwJl450=",
+      "ops/simple_gla/chunk.py": "wmKvR5zeIRZZal2YiPRSXZVp6bJoc5brUfAV1fjbKjw=",
+      "ops/simple_gla/fused_chunk.py": "gN1IUd+GV72DsTw+q9Ye2cMDCrfm9ZYv7dN1/+NHuPs=",
+      "ops/simple_gla/fused_recurrent.py": "UukysuAEhW46kvKCjDNfyfAjNFT2LF4QqZZMvtg+GMc=",
       "ops/simple_gla/naive.py": "g3fAN++vd0aWvZcdOW7Y0boMYcs2mOld8kePWRbl2R0=",
+      "ops/simple_gla/parallel.py": "6jkeEIyN0tHrMGG/PMcPpSHfJK1Gg6OpGpguPhjqulk=",
+      "ops/stickbreaking_attn/__init__.py": "A3p1HEdTYG4f0TLRWGSWQTySjsfOVHbX+Eg7KFyXpQM=",
+      "ops/stickbreaking_attn/naive.py": "P4jkKUVtTg7ImxWWfxbZzFsq//W5jfgeLM7gYVQTClk=",
+      "ops/stickbreaking_attn/parallel.py": "jqY7jl1yRSzCfTO1QtYmXPqKQfS+UdiyVQ1XCFCOKWg=",
       "ops/titans/__init__.py": "4t9OK1aYdW03I+sUlfUjj3ALTez1M/XFD/H6KzeZtGc=",
       "ops/titans/log_impl.py": "Y5tMWJgtVE+GX394DdxQHd6Yirrg3sLx6sinMdQvkAE=",
+      "ops/titans/naive.py": "0Y4R5926mX2CgYI55MjxVfn6p2X4iM4pABCzmVZ9d5c=",
       "ops/ttt/__init__.py": "cFaLRxAWRsKjcmXh01bJbh5U59ZaFaM7DwYgm/KBAbo=",
+      "ops/ttt/chunk.py": "Pyon33URTqVwxnprbp9AC7yZcY99MXBWTOxnf1iAxow=",
+      "ops/ttt/fused_chunk.py": "uVBNMzpMgEozAEuzis/xYVpBrDrxf3mcnWyfw3Lr9Yo=",
+      "ops/ttt/naive.py": "oSWlwYF9ngQShNodpdkaIwb3ndI4RH0jB2jssaKv89s=",
+      "ops/utils/__init__.py": "3NBO3K+KTaQeR4ebY6LYdLOhdUxGFOPSqdAQPLhJt/I=",
+      "ops/utils/cache.py": "qP/aHKYJ7682YgeoetUAIlPTn9S14f2URjRSRdEQ5vk=",
       "ops/utils/constant.py": "uVK9zwOd7kM/DPSQk8cllx5nkYeK0K/KZVyCu0ZFdPg=",
+      "ops/utils/csr.py": "3s/+zgzuhMNIRo9UyriDR+Z94AOfAnx0VAD2VdmCszQ=",
+      "ops/utils/cumsum.py": "0Y9Lu/0qZLxT05txmpYuwvbmuxlgAsF6mnejbOp84xY=",
+      "ops/utils/graph.py": "uj5R0DV8gy3tW/VH9pT4xPn3Zb0CxYJWpH4zaU8KS6E=",
+      "ops/utils/index.py": "KeykCsBnOSNMGGgphH5OKFD/uqiva2rPGKUuqyj/8xk=",
+      "ops/utils/logcumsumexp.py": "ZpYz3i9FZI712XtDqIxCus8A4WamVFlYI8b6tHVbkXQ=",
+      "ops/utils/logsumexp.py": "qklce3jy76dN9WrNT1Bzw5F6Ime+mo00deKYryTMLTk=",
+      "ops/utils/matmul.py": "1X2PdPmFwXEf3ruaLZHBLFrYBwgmh5k+GOP5fuKM5Jw=",
+      "ops/utils/op.py": "Y/Kcks768kujwclindhHtbnQxyJje7mMBVReDvegK1Q=",
+      "ops/utils/pack.py": "92z4AmmkdpJ5rryTNHjB+BkSmlsBU5ssXJqPoLXqdFc=",
+      "ops/utils/pooling.py": "mtioM2xOtxMCEJDBfR+lVxhdg1x3NdBV8PVK+zOb/6Q=",
+      "ops/utils/sinkhorn/__init__.py": "7AY27ZaP2G/gGqo/equFKp9hLHUcAp6Mmi8xE3UGBGc=",
+      "ops/utils/sinkhorn/gluon.py": "bJ9FjAXod5fuweW1T6pP3C1Q6esQRDiS/MchqGZKtIU=",
+      "ops/utils/sinkhorn/naive.py": "WZYZzlJ/atDmXzRnJFFMWfB/t5JIytpzxLz98palzq8=",
+      "ops/utils/sinkhorn/ops.py": "2UH9sL5DsBPRcr+vUgahTvOaI9Myg2cnnZv7oclXx/s=",
+      "ops/utils/softmax.py": "aCTHGnayRqCA6PAd9edtVeb9EoFNfzJdBmDtRZiMZN0=",
+      "ops/utils/softplus.py": "Ykhyi0d+jM2zEyFp8fQqmFNqQmsU5XzUbUpn+jr/Rn0=",
+      "ops/utils/solve_tril.py": "4j2C13HQCr+F5URwwLUO4t3Dag6upnlpGwNew17nphs=",
       "ops/wall_attn/__init__.py": "r/FbDGjGO8WK1DaMbl9axWHpENUa3pzSc0v9qviHJ6I=",
+      "ops/wall_attn/decode.py": "o5q0fsJTiz+0XW0RPFFI1hBXz6PXSBMSTvog8Bc+sGA=",
       "ops/wall_attn/naive.py": "9YKwBSvQ9imy229T5KMYKpuVNb+xEzhXsXdMemRWqjQ=",
+      "ops/wall_attn/parallel.py": "k1DMCeZ/dLyUrpaLt1eWphG3WA0STXIwofUU2/lZx1g=",
+      "symbols.json": "XeDbgUU3taJ6OhnIi2xuHQpY5Guf0MQ5YL5XFWYzxl0=",
+      "utils/__init__.py": "+7iYDYAJEDL9y/TTl/bNr5/O0LpSylNyE6J7c1OCAxs=",
+      "utils/decorators.py": "Xy1DCGkQ7HRyffUHEXsVLhn6LE7xgdIY9Yob1GFm/f8=",
+      "utils/env.py": "LJZbQTALlKpcS2dwt7JBv/AdlZ8Yh7LC7tCDDJs0D2M=",
+      "utils/hardware.py": "2AA4GxTL1WNYIzRJQF3OXCCuATnobx928o6bjXD0U04=",
+      "utils/layers.py": "afB10gGysmmVfVRlPnduPpuvcdAeDAstmqPC0sYJbRw=",
+      "utils/testing.py": "KXWwCDYOxKLEYFxzZgFfh/3kxzMm+12Vq4IsxJaoPlM="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "commit": "faa8674647aa5ebaab5e997cd11e6e3808d090a9",
       "dirty": false
     }
   }

build/torch-cuda/metadata.json.sigstore CHANGED Viewed

@@ -1 +1 @@

- {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrOgAwIBAgIUQ629eN4hUTIWupwbogDv3RA6AyEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMxODIzWhcNMjYxMDE5MTMyODIzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEpl17UO4EUwHQu4gYHIYYLubpp8Zzl5lYFeFTp4LxxeYjaO9yg8RCS6mWZc2nLEcA3HNjkFi4C6oGo/oG5gs1KaOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTU0OTYzNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVFAuVgAABAMASDBGAiEAtnEys30x7Ya9hsVCz7RWdn6nYXSqgMVDe8yHwixSiiYCIQCdcOvWSnZSR+oCBuEqRzBjvDTsmBBOA7bQNQdRM+tFVjAKBggqhkjOPQQDAwNoADBlAjB44h8gf51LU8aLNgDS8WUfSC7PpbdOXA/MDkX3XPLQa7tTAnTW9IgnghlRkSCVIdUCMQDgufgY6hRhjInveoC0gxGEucYH3gdJZ7qJkooniD5e5OD64vpL9D2zpvlaZC+N+Y8="},"tlogEntries":[{"logIndex":"3321382366","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415903","inclusionPromise":{"signedEntryTimestamp":"MEQCIA8Pb2l/03SWHSHoQlVPnuD3qoB1EObRvw8TSvdRmXHCAiBLQW9zMHeh1g3bHLP1y44z5xADQsM8tsSXvzGp0GMBYA=="},"inclusionProof":{"logIndex":"3199478104","rootHash":"4tHHO4qqAF/BE6oVRgF0boQq/yTfvz0KIX8pAsRo3G0=","treeSize":"3199478328","hashes":["boitsHdOY++pKowb8ip8eWyv+giOG6c3LLDf4d59mB0=","qcf970yQgX2TVWwq1E7iCbQrykhEhaftOyMyVe+h/N4=","EslsDaCqZixXslWYzeNOIgKnH6afg5sQlM2aEpV4jq0=","FtY90QTdG+QWhCzKV56/dsHTXqKp04sOsj4tjwrBhI8=","1Du/xr+ro1zVy3zI5m0v/GyxJcdIShmHLA8zKq/JJrc=","6akUiAwNlsFQQrqdMshpY6Nr7bcHUIKiutTPw9YSiOY=","eJJgNMmRIOUXkgOpxCA8t3Iu0nU2BvsYyb5Fuo48Z/s=","gqCLyzcTY8N5+rych9c4HRspuMPAmRkae5U4ACwbpaM=","r4uw63XA1OsTh4/HhkQY/Veqn7iADlv2k7521RDZtlw=","mIeT670PcvMS00qGZ9uqtj9ZBspJzyCniPsKBf9FQ1M=","OXA2SiOGgZ4Atk3yCE3v98reKXHty7rZAo82yZi0yls=","eE+LYckVfbRCjJTwQL99tuBkmhXXjKkrIh1lVDPm+hs=","WF6g44nOcJ5AK5U4U6CvcX9NPmGRLXdvfv+hcf011hw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199478328\n4tHHO4qqAF/BE6oVRgF0boQq/yTfvz0KIX8pAsRo3G0=\n\n— rekor.sigstore.dev wNI9ajBEAiANt/PrcbDQqVmXByPf3zWgItoN7m+bQariPJGNVfk0hwIgNtxR5MYteCwgnROrvIy4TnFG1I1+jhaTEfNh7ViG12U=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwZTE2NWUzMmY5ZDdlNzEyNWVkMjJhYWZlMTY1YzNiYjEyNjA4MThmN2E3YjUwZDdhNTkyOGE4MTI5OWQ1ZjRjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJR3VFSllNbGZJcGZCZXBqMFN1Nk5iLzMxSEdibVpuZzVVRnRUajJwWStPSUFpRUE4R1d4RDlNOXJVdU1od0pMVU53RERFR1ZWVHJ0MXViUHFsOUl1b2R5SDdvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KUFowRjNTVUpCWjBsVlVUWXlPV1ZPTkdoVlZFbFhkWEIzWW05blJIWXpVa0UyUVhsRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTRUMFJKZWxkb1kwNU5hbGw0VFVSRk5VMVVUWGxQUkVsNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ3YkRFM1ZVODBSVlYzU0ZGMU5HZFpTRWxaV1V4MVluQndPRnA2YkRWc1dVWmxSbFFLY0RSTWVIaGxXV3BoVHpsNVp6aFNRMU0yYlZkYVl6SnVURVZqUVROSVRtcHJSbWswUXpadlIyOHZiMGMxWjNNeFMyRlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVVlRCUFZGbDZUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtaQmRWWm5RVUZDUVUxQlUwUkNSMEZwUlVGMGJrVjVjek13ZURkWllUbG9jMVpEZWpkU1YyUnVObTVaV0ZOeFowMVdSQXBsT0hsSWQybDRVMmxwV1VOSlVVTmtZMDkyVjFOdVdsTlNLMjlEUW5WRmNWSjZRbXAyUkZSemJVSkNUMEUzWWxGT1VXUlNUU3QwUmxacVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2IwRkVRbXhCYWtJME5HZzRaMlkxTVV4Vk9HRk1UbWRFVXpoWFZXWlRRemRRY0dKa1QxaEJMMDFFYTFneldGQk1VV0UzZEZRS1FXNVVWemxKWjI1bmFHeFNhMU5EVmtsa1ZVTk5VVVJuZFdabldUWm9VbWhxU1c1MlpXOURNR2Q0UjBWMVkxbElNMmRrU2xvM2NVcHJiMjl1YVVRMVpRbzFUMFEyTkhad1REbEVNbnB3ZG14aFdrTXJUaXRaT0QwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgM+jtpRT7VTayyXNU6XeuWevLOHSgymjIR3AhSqAF5ZoCFDPzZCiwDE+g/BiHrxxd5OfIKibdGA8yMDI2MTAxOTEzMTgyM1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMxODIzWjAvBgkqhkiG9w0BCQQxIgQgChdtG/mcWGdAqUjlHsL98rTG5klQTLXZofRCueCmlcMwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxANvTfbN0SFidKmyitQtHD4Zqb8S9dKWiNCxfM34Oh3cGtp11St3BpK8lBFE6kVV5ygIxAPF6xXECkECEvGUpxdM+YCyEz16v/A6sZHjToLqQ52fY6uNObrkl1Lyv7ghcHDtBQQ=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"DhZeMvnX5xJe0iqv4WXDuxJggY96e1DXpZKKgSmdX0w="},"signature":"MEUCIGuEJYMlfIpfBepj0Su6Nb/31HGbmZng5UFtTj2pY+OIAiEA8GWxD9M9rUuMhwJLUNwDDEGVVTrt1ubPql9IuodyH7o="}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUW8juOVkziLnnecCKBESbrqbEgJEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDIzMTEzNjAyWhcNMjYxMDIzMTE0NjAyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEnvx47gGnlRzCN/4OkJ0a7wqYEsIFmCINUHZaFbz5VXCd4BWhf0QldM2OnXd4Ogo0yCWW1dQCX18MRUBj9CRGX6OCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmYWE4Njc0NjQ3YWE1ZWJhYWI1ZTk5N2NkMTFlNmUzODA4ZDA5MGE5MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmYWE4Njc0NjQ3YWE1ZWJhYWI1ZTk5N2NkMTFlNmUzODA4ZDA5MGE5MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmYWE4Njc0NjQ3YWE1ZWJhYWI1ZTk5N2NkMTFlNmUzODA4ZDA5MGE5MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmFhODY3NDY0N2FhNWViYWFiNWU5OTdjZDExZTZlMzgwOGQwOTBhOTAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTk0NjY4MTMwNi9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhaIvpVQAABAMARzBFAiAoSbfUgNXfezTQtR5qo/CAMwLW8dHV4YbAmcGsMcaFGAIhAMOHSfyt8UUJmPp+JbjTW2xG1BtB/lVIRWOl5oJ/2vZMMAoGCCqGSM49BAMDA2kAMGYCMQCh4ra7/gf8kr1+MdulHjJAXsUd6o0Y9UypE3gFlmHTIgsa/PX9g9NIQP2SRERurOICMQDs43kZxEJ3O3bKz90jbuYM8eNXViJ/TqCHiYc2lWLDQ3rLJWze8xhx7Ks+OwnIgBQ="}, "tlogEntries":[{"logIndex":"3419000459", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792755362", "inclusionPromise":{"signedEntryTimestamp":"MEUCICKD7ZdvH1zemVgy07V5zNzmAv0zBIb/8XOEU9QuaaUEAiEAlc/U/9oh4oaQkyOWIzDOALZtdsEroH6lxTDoz9oWPSE="}, "inclusionProof":{"logIndex":"3297096197", "rootHash":"JEwIlEh4C2tvEMwtPOSql67Ts10gAHO7oKqbcfoRsd8=", "treeSize":"3297096204", "hashes":["eKPdjKmW06i/+OgD5Op1qaQr5ljD9icDNR9hKFkKoGs=", "7rIrg/UOeaKdEevZpZHUL3OkluOm/jIeaO923YSkHYo=", "DFFxkRWq9l1mck3McpCgX5nsFudwuRn4U349BXtS7Po=", "S1S/TJWDErr9PeC72czEVORauVEpucJfdtz2vOI5UII=", "06+aOxxaVl4XdwkhaYb3Dh+97wX+oKPonkJn+6nGgmE=", "4iY88O6LVkiDY854scCm32PtSUJYHIsmX33LyfcvBNs=", "DAFTN+3Ntn00j3uqF6DuxpoDvddLg+hpOGsbC8FJr/g=", "sH3iRJ+YSpMgcl9MQvq5mOnxKnX0UOS6Qxva+JrjED8=", "mIjtw1KA13/4jqXfKlJiFrODS7dNF+hqEUNlXprnzTE=", "pNenO6BoAxc3BwcFEO9IJI4uZmwdkynjP9MI8OiYc50=", "1GPTGQWDthSzUsPF7t/XQmpI2SH4KZm3gaVrhMs6T2k=", "YYCr61A8azar3dfO8Ct+pri7LqDxkUlUkfs7fiqYErU=", "1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3297096204\nJEwIlEh4C2tvEMwtPOSql67Ts10gAHO7oKqbcfoRsd8=\n\n— rekor.sigstore.dev wNI9ajBFAiEA3J9FLPnuB38IvvWseghTN705Z1/I82uJ9pjXiomV8lcCIF/CifFOxrfvDzHBxDm21wUegDzbni28WhNQ9fftJoeE\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwZGIyMmFhMzI4ZWUzZjcwOTAwYTFmMTM4MzRmYjFiMmFiNDE5NTMwODY0N2MzNWM5NjFmMTAzMGQyNzljMGJhIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURDL2YydTZCS013bFJ5U3Z0RWdUdUsrbUVjTGFQT1lYNmRYUkFwNWdSMytRSWhBS1p3bXRKZ2d2TDNVMEF3SU5kM3NWYXRsN2xCK3oxRjY3V0ptcm1IOVpiQSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlZ6aHFkVTlXYTNwcFRHNXVaV05EUzBKRlUySnljV0pGWjBwRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1hwTlZFVjZUbXBCZVZkb1kwNU5hbGw0VFVSSmVrMVVSVEJPYWtGNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1ZG5nME4yZEhibXhTZWtOT0x6UlBhMG93WVRkM2NWbEZjMGxHYlVOSlRsVklXbUVLUm1KNk5WWllRMlEwUWxkb1pqQlJiR1JOTWs5dVdHUTBUMmR2TUhsRFYxY3haRkZEV0RFNFRWSlZRbW81UTFKSFdEWlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V1ZkRk5BcE9hbU13VG1wUk0xbFhSVEZhVjBwb1dWZEpNVnBVYXpWT01rNXJUVlJHYkU1dFZYcFBSRUUwV2tSQk5VMUhSVFZOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxWlYwVTBUbXBqTUU1cVVUTlpWMFV4Q2xwWFNtaFpWMGt4V2xSck5VNHlUbXROVkVac1RtMVZlazlFUVRSYVJFRTFUVWRGTlUxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxWlYwVTBDazVxWXpCT2FsRXpXVmRGTVZwWFNtaFpWMGt4V2xSck5VNHlUbXROVkVac1RtMVZlazlFUVRSYVJFRTFUVWRGTlUxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlVab1QwUlpNMDVFV1RCT01rWm9UbGRXYVZsWFJtbE9WMVUxVDFSa2FscEVSWGhhVkZwc1RYcG5kMDlIVVhjS1QxUkNhRTlVUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUnJNRTVxV1RSTlZFMTNUbWs1YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvWVVsMmNGWlJRVUZDUVUxQlVucENSa0ZwUVc5VFltWlZaMDVZWm1WNlZGRjBValZ4Ynk5RFFVMTNURmM0WkVoV05GbGlRUXB0WTBkelRXTmhSa2RCU1doQlRVOUlVMlo1ZERoVlZVcHRVSEFyU21KcVZGY3llRWN4UW5SQ0wyeFdTVkpYVDJ3MWIwb3ZNblphVFUxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVOb05ISmhOeTluWmpocmNqRXJUV1IxYkVocVNrRlljMVZrTm04d1dUbFZlWEJGTTJkR2JHMUlWRWxuYzJFS0wxQllPV2M1VGtsUlVESlRVa1ZTZFhKUFNVTk5VVVJ6TkROclduaEZTak5QTTJKTGVqa3dhbUoxV1UwNFpVNVlWbWxLTDFSeFEwaHBXV015YkZkTVJBcFJNM0pNU2xkNlpUaDRhSGczUzNNclQzZHVTV2RDVVQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg3bpS2VtJbbj+NwI4jCmvshiu/17s+rjceNrg2WK85ewCFQC8epRj5pgyFllwm2wpZ1CZZQQv6hgPMjAyNjEwMjMxMTM2MDJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAyMzExMzYwMlowLwYJKoZIhvcNAQkEMSIEICpkIWImbVOed+fyc52Sr18JzDELBbbea0FI1ljUg31hMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQD2N3QvdKUrb7mICl3pBDsE03F8HuaqvKzbzOufuc/onqDuhozEGeCO7WtAfnBimSYCMQDaZJiF4IvNVlXG7SnHGbx6kg6iHsKj9VRRlszRA/TNd923wrJtH7f1Kck8QNXzJgU="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"DbIqoyjuP3CQCh8Tg0+xsqtBlTCGR8Nclh8QMNJ5wLo="}, "signature":"MEYCIQDC/f2u6BKMwlRySvtEgTuK+mEcLaPOYX6dXRAp5gR3+QIhAKZwmtJggvL3U0AwINd3sVatl7lB+z1F67WJmrmH9ZbA"}}

build/torch-cuda/modules/__init__.py CHANGED Viewed

@@ -5,25 +5,106 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-from ..modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution
 from ..modules.fused_bitlinear import BitLinear, FusedBitLinear
 from ..modules.fused_cross_entropy import FusedCrossEntropyLoss
 from ..modules.fused_kl_div import FusedKLDivLoss
 from ..modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
-from ..modules.fused_norm_gate import (
     FusedLayerNormGated,
     FusedLayerNormSwishGate,
     FusedLayerNormSwishGateLinear,
     FusedRMSNormGated,
     FusedRMSNormSwishGate,
     FusedRMSNormSwishGateLinear,
 )
-from ..modules.l2norm import L2Norm
-from ..modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear
-from ..modules.mlp import GatedMLP
 from ..modules.rotary import RotaryEmbedding
 from ..modules.token_shift import TokenShift
 __all__ = [
     'BitLinear',
     'FusedBitLinear',

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+import sys
+from importlib import import_module
+from importlib.machinery import ModuleSpec
+from types import ModuleType
+from ..modules.causal_conv1d import ShortConvolution
 from ..modules.fused_bitlinear import BitLinear, FusedBitLinear
 from ..modules.fused_cross_entropy import FusedCrossEntropyLoss
 from ..modules.fused_kl_div import FusedKLDivLoss
 from ..modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
+from ..modules.long_conv import ImplicitLongConvolution, LongConvolution
+from ..modules.mlp import GatedMLP
+from ..modules.norm import (
     FusedLayerNormGated,
     FusedLayerNormSwishGate,
     FusedLayerNormSwishGateLinear,
     FusedRMSNormGated,
     FusedRMSNormSwishGate,
     FusedRMSNormSwishGateLinear,
+    GroupNorm,
+    GroupNormLinear,
+    L2Norm,
+    LayerNorm,
+    LayerNormLinear,
+    RMSNorm,
+    RMSNormLinear,
 )
 from ..modules.rotary import RotaryEmbedding
 from ..modules.token_shift import TokenShift
+_MODULE_ALIASES = {
+    'fused_norm_gate': ('norm.fused_norm_gate',),
+    'l2norm': ('norm.l2norm',),
+    'layernorm': ('norm.layernorm',),
+    'layernorm_gated': ('norm.layernorm_gated',),
+    'convolution': ('causal_conv1d', 'long_conv', 'causal_conv1d.cp', 'causal_conv1d.backends.cuda'),
+    'conv': ('causal_conv1d', 'long_conv'),
+    'conv.causal_conv1d': ('causal_conv1d.ops',),
+    'conv.short_conv': ('causal_conv1d.ops',),
+    'conv.long_conv': ('long_conv',),
+    'conv.cp': ('causal_conv1d.cp',),
+    'conv.cp.ops': ('causal_conv1d.cp',),
+    'conv.cuda': ('causal_conv1d.backends.cuda',),
+    'conv.cuda.ops': ('causal_conv1d.backends.cuda.ops',),
+    'conv.triton': ('causal_conv1d.ops',),
+    'conv.triton.ops': ('causal_conv1d.ops',),
+    'conv.triton.kernels': ('causal_conv1d.ops',),
+}
+for _old, _targets in _MODULE_ALIASES.items():
+    _fullname = f'{__name__}.{_old}'
+    if len(_targets) == 1:
+        _module = import_module(f'{__name__}.{_targets[0]}')
+    else:
+        _module = ModuleType(_fullname)
+        _module.__spec__ = ModuleSpec(_fullname, loader=None, is_package=True)
+        _module.__path__ = []
+        for _target in _targets:
+            _source = import_module(f'{__name__}.{_target}')
+            for _symbol in getattr(_source, '__all__', vars(_source)):
+                if not _symbol.startswith('_'):
+                    vars(_module).setdefault(_symbol, getattr(_source, _symbol))
+        _module.__all__ = [name for name in vars(_module) if not name.startswith('_')]
+    sys.modules[_fullname] = _module
+    # conv.causal_conv1d is a function on the old package.
+    if _old != 'conv.causal_conv1d':
+        _parent, _, _child = _fullname.rpartition('.')
+        setattr(sys.modules[_parent], _child, _module)
+# defer GRPO's compile policy until it is explicitly imported
+for _name in ('activations', 'rotary', 'fused_cross_entropy', 'fused_kl_div', 'fused_linear_cross_entropy'):
+    _module = import_module(f'{__name__}.{_name}')
+    for _symbol, _value in vars(import_module(f'{__name__}.{_name}.ops')).items():
+        if not _symbol.startswith('_'):
+            vars(_module).setdefault(_symbol, _value)
+_SYMBOL_ALIASES = {
+    'fused_bitlinear': {
+        'layer_norm_fwd_quant': 'norm.layernorm_quant.layer_norm_quant_fwd',
+        'layer_norm_bwd': 'norm.layernorm_quant.layer_norm_quant_bwd',
+        'LayerNormLinearQuantFn': 'LayerNormLinearQuantFunction',
+        'layer_norm_linear_quant_fn': 'layer_norm_linear_quant',
+    },
+    'fused_cross_entropy': {
+        'fused_cross_entropy_forward': 'cross_entropy_fwd',
+        'CrossEntropyLossFunction': 'FusedCrossEntropyFunction',
+    },
+    'fused_kl_div': {'fused_kl_div_forward': 'fused_kl_div_fwd', 'fused_kl_div_backward': 'fused_kl_div_bwd'},
+    'fused_linear_cross_entropy': {
+        'fused_linear_cross_entropy_forward': 'fused_linear_cross_entropy_fwd',
+        'fused_linear_cross_entropy_backward': 'fused_linear_cross_entropy_bwd',
+    },
+}
+for _name, _aliases in _SYMBOL_ALIASES.items():
+    _module = sys.modules[f'{__name__}.{_name}']
+    for _old, _new in _aliases.items():
+        _target, _, _symbol = _new.rpartition('.')
+        _source = import_module(f'{__name__}.{_target}') if _target else _module
+        setattr(_module, _old, getattr(_source, _symbol))
 __all__ = [
     'BitLinear',
     'FusedBitLinear',

build/{torch-xpu/ops/cp → torch-cuda/modules/activations}/__init__.py RENAMED Viewed

@@ -5,28 +5,35 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-# Context Parallel operators and utilities
-from .comm import (
-    all_gather_into_tensor,
-    all_reduce_sum,
-    conv_cp_send_recv_bwd,
-    conv_cp_send_recv_fwd,
-    send_recv_bwd,
-    send_recv_fwd,
-)
-from .context import (
-    FLACPContext,
-    build_cp_context,
 )
 __all__ = [
-    "FLACPContext",
-    "all_gather_into_tensor",
-    "all_reduce_sum",
-    "build_cp_context",
-    "conv_cp_send_recv_bwd",
-    "conv_cp_send_recv_fwd",
-    "send_recv_bwd",
-    "send_recv_fwd",
 ]

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+from ...modules.activations.ops import (
+    ACT2FN,
+    elu_p1,
+    fast_gelu_impl,
+    logsigmoid,
+    powglu,
+    powglu_linear,
+    sigmoid,
+    sigmoidglu,
+    sigmoidglu_linear,
+    sqrelu,
+    swiglu,
+    swiglu_linear,
+    swish,
 )
 __all__ = [
+    'ACT2FN',
+    'elu_p1',
+    'fast_gelu_impl',
+    'logsigmoid',
+    'powglu',
+    'powglu_linear',
+    'sigmoid',
+    'sigmoidglu',
+    'sigmoidglu_linear',
+    'sqrelu',
+    'swiglu',
+    'swiglu_linear',
+    'swish',
 ]

build/torch-cuda/modules/{activations.py → activations/ops.py} RENAMED Viewed

@@ -8,43 +8,58 @@
 """Fused activation kernels.
 The kernels address their inputs through the row stride instead of assuming a fully contiguous buffer.
-An inner-contiguous input — such as one half of ``x.chunk(2, dim=-1)`` — is therefore read in place, sparing the extra
-``.contiguous()`` copy (and its memory traffic) that a plain flat element-wise kernel would force on every call.
 """
 import torch
 import torch.nn.functional as F
 import triton
 import triton.language as tl
-from ..modules.backends import dispatch
-from ..ops.utils.op import exp, log
-from ..utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard
-NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32]
 def _get_stride(x: torch.Tensor) -> int:
     """Get the row stride for viewing a tensor as 2D (num_rows, D) where D = shape[-1].
     Returns stride(-2) if the tensor is at least 2D, or 0 for 1D tensors.
-    The caller must ensure the tensor is "inner-contiguous" (stride(-1) == 1 and
-    higher dims are contiguous relative to dim -2) before using this value.
     """
     if x.ndim < 2:
         return 0
     return x.stride(-2)
 def _is_inner_contiguous(x: torch.Tensor) -> bool:
     """Check if a tensor can be safely viewed as 2D (num_rows, D) with row stride = stride(-2).
-    This holds when stride(-1) == 1 and all dimensions above -2 are contiguous
-    with respect to the dimension below them.
     """
     ndim = x.ndim
     if ndim < 2:
-        return True
     if x.stride(-1) != 1:
         return False
     if ndim == 2:
@@ -69,7 +84,9 @@ def _is_inner_contiguous(x: torch.Tensor) -> bool:
 def _ensure_inner_contiguous(x: torch.Tensor) -> torch.Tensor:
     """Make the tensor inner-contiguous if it isn't already."""
-    if _is_inner_contiguous(x):
         return x
     return x.contiguous()
@@ -77,32 +94,17 @@ def _ensure_inner_contiguous(x: torch.Tensor) -> torch.Tensor:
 def _alloc_output(x: torch.Tensor, contiguous: bool = False) -> torch.Tensor:
     """Allocate the output: a fresh contiguous buffer, or ``empty_like`` otherwise.
-    ``empty_like`` keeps the input's memory format only when it is dense; a non-dense
-    strided view (e.g. a ``chunk`` slice) falls back to contiguous, not the input stride.
     """
     if contiguous:
         return x.new_empty(x.shape)
     return torch.empty_like(x)
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def sigmoid_fwd_kernel(
-    x, y,
-    stride_x_row,
-    stride_y_row,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -113,25 +115,9 @@ def sigmoid_fwd_kernel(
     tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def sigmoid_bwd_kernel(
-    x, dy, dx,
-    stride_x_row,
-    stride_dy_row,
-    stride_dx_row,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -144,35 +130,35 @@ def sigmoid_bwd_kernel(
     tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
-@dispatch('modules')
 def sigmoid_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
     T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x, output_contiguous)
     sigmoid_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
         T=T,
         D=D,
     )
     return y
-@dispatch('modules')
 def sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    dy = _ensure_inner_contiguous(dy)
     T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
     sigmoid_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         dy=dy,
         dx=dx,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(dy),
-        stride_dx_row=_get_stride(dx),
         T=T,
         D=D,
     )
@@ -185,38 +171,111 @@ class SigmoidFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x):
         ctx.save_for_backward(x)
-        return sigmoid_fwd(x)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, = ctx.saved_tensors
-        return sigmoid_bwd(x, dout)
 sigmoid = SigmoidFunction.apply
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def logsigmoid_fwd_kernel(
-    x,
-    y,
-    stride_x_row,
-    stride_y_row,
-    temperature,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -229,15 +288,7 @@ def logsigmoid_fwd_kernel(
     tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
 def logsigmoid_bwd_kernel(
     x,
@@ -262,16 +313,16 @@ def logsigmoid_bwd_kernel(
     tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
-@dispatch('modules')
 def logsigmoid_fwd(x: torch.Tensor, temperature: float = 1., output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
     T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x, output_contiguous)
     logsigmoid_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
         temperature=temperature,
         T=T,
         D=D,
@@ -279,24 +330,24 @@ def logsigmoid_fwd(x: torch.Tensor, temperature: float = 1., output_contiguous:
     return y
-@dispatch('modules')
 def logsigmoid_bwd(
     x: torch.Tensor,
     dy: torch.Tensor,
     temperature: float = 1.,
     output_contiguous: bool = False,
 ) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    dy = _ensure_inner_contiguous(dy)
     T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
     logsigmoid_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         dy=dy,
         dx=dx,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(dy),
-        stride_dx_row=_get_stride(dx),
         temperature=temperature,
         T=T,
         D=D,
@@ -311,37 +362,22 @@ class LogSigmoidFunction(torch.autograd.Function):
     def forward(ctx, x, temperature):
         ctx.save_for_backward(x)
         ctx.temperature = temperature
-        return logsigmoid_fwd(x, temperature)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dy):
         x, = ctx.saved_tensors
-        return logsigmoid_bwd(x, dy, ctx.temperature), None
 def logsigmoid(x: torch.Tensor, temperature: float = 1.) -> torch.Tensor:
     return LogSigmoidFunction.apply(x, temperature)
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def swish_fwd_kernel(
-    x, y,
-    stride_x_row,
-    stride_y_row,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -352,25 +388,9 @@ def swish_fwd_kernel(
     tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def swish_bwd_kernel(
-    x, dy, dx,
-    stride_x_row,
-    stride_dy_row,
-    stride_dx_row,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -383,35 +403,35 @@ def swish_bwd_kernel(
     tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
-@dispatch('modules')
 def swish_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
     T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x, output_contiguous)
     swish_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
         T=T,
         D=D,
     )
     return y
-@dispatch('modules')
 def swish_bwd(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    dy = _ensure_inner_contiguous(dy)
     T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
     swish_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         dy=dy,
         dx=dx,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(dy),
-        stride_dx_row=_get_stride(dx),
         T=T,
         D=D,
     )
@@ -424,13 +444,13 @@ class SwishFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x):
         ctx.save_for_backward(x)
-        return swish_fwd(x)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, = ctx.saved_tensors
-        return swish_bwd(x, dout)
 swish = SwishFunction.apply
@@ -440,27 +460,20 @@ swish = SwishFunction.apply
 # sqrt(2/pi)  -> 0.79788456
-# this function is tanh approximation of gelu
-# actual gelu is:
-# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
 @torch.compile
 def bias_gelu(y, bias):
     x = bias + y
     return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype)
-# gradient of tanh approximation of gelu
-# gradient of actual gelu is:
-# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
 @torch.compile
 def bias_gelu_bwd(g, y, bias):
-    """Assume that y has shape (B, D=D) and bias has shape (D)"""
     x = bias + y
     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
-    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (
-        1 + tanh_out
-    )
     grad_y = ff * g
     return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype)
@@ -468,54 +481,43 @@ def bias_gelu_bwd(g, y, bias):
 class GeLUFunction(torch.autograd.Function):
     @staticmethod
-    # bias is an optional argument
     def forward(ctx, input, bias):
         ctx.save_for_backward(input, bias)
-        return bias_gelu(input, bias)
     @staticmethod
     def backward(ctx, grad_output):
         input, bias = ctx.saved_tensors
-        return bias_gelu_bwd(grad_output, input, bias)
 bias_gelu_impl = GeLUFunction.apply
-# this function is tanh approximation of gelu
-# actual gelu is:
-# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
-@dispatch('modules')
 @torch.compile
 def gelu_fwd(x):
     return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype)
-# gradient of tanh approximation of gelu
-# gradient of actual gelu is:
-# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
-@dispatch('modules')
 @torch.compile
 def gelu_bwd(g, x):
     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
-    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (
-        1 + tanh_out
-    )
     return (ff * g).to(dtype=x.dtype)
 class FastGeLUFunction(torch.autograd.Function):
     @staticmethod
-    # bias is an optional argument
     def forward(ctx, input):
         ctx.save_for_backward(input)
-        return gelu_fwd(input)
     @staticmethod
     def backward(ctx, grad_output):
         (input,) = ctx.saved_tensors
-        tmp = gelu_bwd(grad_output, input)
         return tmp
@@ -527,14 +529,12 @@ def relu_bwd(g, x):
     return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype)
-@dispatch('modules')
 @torch.compile
 def sqrelu_fwd(x):
     r = F.relu(x.float())
     return (r * r).to(dtype=x.dtype)
-@dispatch('modules')
 @torch.compile
 def sqrelu_bwd(g, x):
     return (2.0 * g * F.relu(x.float())).to(dtype=x.dtype)
@@ -545,36 +545,20 @@ class SquaredReLUFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input):
         ctx.save_for_backward(input)
-        return sqrelu_fwd(input)
     @staticmethod
     def backward(ctx, grad_output):
         input, = ctx.saved_tensors
-        return sqrelu_bwd(grad_output, input)
 sqrelu = SquaredReLUFunction.apply
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def swiglu_fwd_kernel(
-    x, y, z,
-    stride_x_row,
-    stride_y_row,
-    stride_z_row,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -586,21 +570,16 @@ def swiglu_fwd_kernel(
     tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['z'] is not None,
-})
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
 def swiglu_fwdbwd_kernel(
-    x, y, g, dx, dy, z,
     stride_x_row,
     stride_y_row,
     stride_g_row,
@@ -633,27 +612,27 @@ def swiglu_fwdbwd_kernel(
         tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
-@dispatch('modules')
 def swiglu_fwd(x: torch.Tensor, y: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"swiglu_fwd: shape mismatch x={x.shape} y={y.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
     T, D = x.numel(), x.shape[-1]
-    z = _alloc_output(x, output_contiguous)
     swiglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
         z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_z_row=_get_stride(z),
         T=T,
         D=D,
     )
     return z
-@dispatch('modules')
 def swiglu_fwdbwd(
     x: torch.Tensor,
     y: torch.Tensor,
@@ -662,14 +641,14 @@ def swiglu_fwdbwd(
     output_contiguous: bool = False,
 ):
     assert x.shape == y.shape == g.shape, f"swiglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    g = _ensure_inner_contiguous(g)
     T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    dy = _alloc_output(y, output_contiguous)
     if use_weight:
-        z = _alloc_output(x, output_contiguous)
     else:
         z = None
     swiglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
@@ -679,12 +658,12 @@ def swiglu_fwdbwd(
         dx=dx,
         dy=dy,
         z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_g_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        stride_dy_row=_get_stride(dy),
-        stride_z_row=_get_stride(z) if z is not None else 0,
         T=T,
         D=D,
     )
@@ -705,13 +684,13 @@ class SwiGLUFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x, y):
         ctx.save_for_backward(x, y)
-        return swiglu_fwd(x, y)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, y = ctx.saved_tensors
-        return swiglu_fwdbwd(x, y, dout)
 class SwiGLULinearFunction(torch.autograd.Function):
@@ -728,10 +707,9 @@ class SwiGLULinearFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     @autocast_custom_fwd
     def forward(ctx, x, y, weight, bias):
-        z = swiglu_fwd(x, y, output_contiguous=True)
         out = F.linear(z, weight, bias)
         ctx.save_for_backward(x, y, weight)
-        ctx.linear_bias_is_none = bias is None
         return out
     @staticmethod
@@ -740,40 +718,32 @@ class SwiGLULinearFunction(torch.autograd.Function):
     def backward(ctx, dout, *args):
         x, y, weight = ctx.saved_tensors
         dout = dout.reshape(-1, dout.shape[-1])
-        dz = F.linear(dout, weight.t()).view_as(x)
-        dx, dy, z = swiglu_fwdbwd(x, y, dz, use_weight=True, output_contiguous=True)
-        dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1]))
-        dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
         return dx, dy, dlinear_weight, dlinear_bias
 swiglu = SwiGLUFunction.apply
-@dispatch('modules')
 def swiglu_linear(x, y, weight, bias):
     return SwiGLULinearFunction.apply(x, y, weight, bias)
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def sigmoidglu_fwd_kernel(
-    x, y, z,
-    stride_x_row,
-    stride_y_row,
-    stride_z_row,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -785,21 +755,16 @@ def sigmoidglu_fwd_kernel(
     tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['z'] is not None,
-})
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
 def sigmoidglu_fwdbwd_kernel(
-    x, y, g, dx, dy, z,
     stride_x_row,
     stride_y_row,
     stride_g_row,
@@ -834,17 +799,17 @@ def sigmoidglu_fwdbwd_kernel(
 @torch.compiler.disable
 def sigmoidglu_fwd(x: torch.Tensor, y: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"sigmoidglu_fwd: shape mismatch x={x.shape} y={y.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
     T, D = x.numel(), x.shape[-1]
-    z = _alloc_output(x, output_contiguous)
     sigmoidglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
         z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_z_row=_get_stride(z),
         T=T,
         D=D,
     )
@@ -860,14 +825,14 @@ def sigmoidglu_fwdbwd(
     output_contiguous: bool = False,
 ):
     assert x.shape == y.shape == g.shape, f"sigmoidglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    g = _ensure_inner_contiguous(g)
     T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    dy = _alloc_output(y, output_contiguous)
     if use_weight:
-        z = _alloc_output(x, output_contiguous)
     else:
         z = None
     sigmoidglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
@@ -877,12 +842,12 @@ def sigmoidglu_fwdbwd(
         dx=dx,
         dy=dy,
         z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_g_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        stride_dy_row=_get_stride(dy),
-        stride_z_row=_get_stride(z) if z is not None else 0,
         T=T,
         D=D,
     )
@@ -903,13 +868,13 @@ class SigmoidGLUFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x, y):
         ctx.save_for_backward(x, y)
-        return sigmoidglu_fwd(x, y)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, y = ctx.saved_tensors
-        return sigmoidglu_fwdbwd(x, y, dout)
 class SigmoidGLULinearFunction(torch.autograd.Function):
@@ -926,7 +891,7 @@ class SigmoidGLULinearFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     @autocast_custom_fwd
     def forward(ctx, x, y, weight, bias):
-        z = sigmoidglu_fwd(x, y, output_contiguous=True)
         out = F.linear(z, weight, bias)
         ctx.save_for_backward(x, y, weight)
         ctx.linear_bias_is_none = bias is None
@@ -939,7 +904,7 @@ class SigmoidGLULinearFunction(torch.autograd.Function):
         x, y, weight = ctx.saved_tensors
         dout = dout.reshape(-1, dout.shape[-1])
         dz = F.linear(dout, weight.t()).view_as(x)
-        dx, dy, z = sigmoidglu_fwdbwd(x, y, dz, use_weight=True, output_contiguous=True)
         dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1]))
         dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
         return dx, dy, dlinear_weight, dlinear_bias
@@ -951,26 +916,9 @@ sigmoidglu = SigmoidGLUFunction.apply
 sigmoidglu_linear = SigmoidGLULinearFunction.apply
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
-def powglu_fwd_kernel(
-    x, y, z,
-    stride_x_row,
-    stride_y_row,
-    stride_z_row,
-    m,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
@@ -990,21 +938,16 @@ def powglu_fwd_kernel(
     tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['z'] is not None,
-})
-@triton.autotune(
-    configs=[
-        triton.Config({'B': bs}, num_warps=num_warps)
-        for bs in [512, 1024, 2048, 4096, 8192]
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=['D'],
-    **autotune_cache_kwargs,
-)
 @triton.jit(do_not_specialize=['T'])
 def powglu_fwdbwd_kernel(
-    x, y, g, dx, dy, z,
     stride_x_row,
     stride_y_row,
     stride_g_row,
@@ -1054,20 +997,20 @@ def powglu_fwdbwd_kernel(
         tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
-@dispatch('modules')
 def powglu_fwd(x: torch.Tensor, y: torch.Tensor, power: float = 3.0, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"powglu_fwd: shape mismatch x={x.shape} y={y.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
     T, D = x.numel(), x.shape[-1]
-    z = _alloc_output(x, output_contiguous)
     powglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
         z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_z_row=_get_stride(z),
         m=power,
         T=T,
         D=D,
@@ -1075,7 +1018,7 @@ def powglu_fwd(x: torch.Tensor, y: torch.Tensor, power: float = 3.0, output_cont
     return z
-@dispatch('modules')
 def powglu_fwdbwd(
     x: torch.Tensor,
     y: torch.Tensor,
@@ -1085,14 +1028,14 @@ def powglu_fwdbwd(
     output_contiguous: bool = False,
 ):
     assert x.shape == y.shape == g.shape, f"powglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    g = _ensure_inner_contiguous(g)
     T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    dy = _alloc_output(y, output_contiguous)
     if use_weight:
-        z = _alloc_output(x, output_contiguous)
     else:
         z = None
     powglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
@@ -1102,12 +1045,12 @@ def powglu_fwdbwd(
         dx=dx,
         dy=dy,
         z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_g_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        stride_dy_row=_get_stride(dy),
-        stride_z_row=_get_stride(z) if z is not None else 0,
         m=power,
         T=T,
         D=D,
@@ -1125,8 +1068,8 @@ class PowGLUFunction(torch.autograd.Function):
         \text{PowGLU}(x, y) = g(x) * y,\quad
         g(x) = \begin{cases} x^{power/(\sqrt{x}+1)}\,\sigma(x) & x > 0 \\ x\,\sigma(x) & x \le 0 \end{cases}
-    For ``x <= 0`` the gate reduces to swish, matching SwiGLU; for large ``x > 0`` it saturates instead of
-    growing, replacing SwiGLU's quadratic amplification with bounded growth (Power Linear Unit, arXiv:2605.25704).
     """
     @staticmethod
@@ -1134,13 +1077,13 @@ class PowGLUFunction(torch.autograd.Function):
     def forward(ctx, x, y, power):
         ctx.save_for_backward(x, y)
         ctx.power = power
-        return powglu_fwd(x, y, power)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, y = ctx.saved_tensors
-        dx, dy = powglu_fwdbwd(x, y, dout, ctx.power)
         return dx, dy, None
@@ -1158,7 +1101,7 @@ class PowGLULinearFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     @autocast_custom_fwd
     def forward(ctx, x, y, weight, bias, power):
-        z = powglu_fwd(x, y, power, output_contiguous=True)
         out = F.linear(z, weight, bias)
         ctx.save_for_backward(x, y, weight)
         ctx.linear_bias_is_none = bias is None
@@ -1172,7 +1115,7 @@ class PowGLULinearFunction(torch.autograd.Function):
         x, y, weight = ctx.saved_tensors
         dout = dout.reshape(-1, dout.shape[-1])
         dz = F.linear(dout, weight.t()).view_as(x)
-        dx, dy, z = powglu_fwdbwd(x, y, dz, ctx.power, use_weight=True, output_contiguous=True)
         dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1]))
         dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
         return dx, dy, dlinear_weight, dlinear_bias, None
@@ -1182,7 +1125,7 @@ def powglu(x: torch.Tensor, y: torch.Tensor, power: float = 3.0) -> torch.Tensor
     return PowGLUFunction.apply(x, y, power)
-@dispatch('modules')
 def powglu_linear(
     x: torch.Tensor,
     y: torch.Tensor,

 """Fused activation kernels.
 The kernels address their inputs through the row stride instead of assuming a fully contiguous buffer.
+An inner-contiguous input — such as one half of ``x.chunk(2, dim=-1)`` — is therefore read in place,
+avoiding the ``.contiguous()`` copy and memory traffic that a flat element-wise kernel would require on every call.
 """
 import torch
 import torch.nn.functional as F
 import triton
 import triton.language as tl
+from triton.language.extra import libdevice
+from ...backends import dispatch
+from ...ops.utils.op import exp, log
+from ...utils import IS_AMD, IS_INTEL, IS_NPU, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard
+def _activation_autotune_configs():
+    """Return autotune configs scaled to the device's scratch-space budget."""
+    if IS_INTEL:
+        # Intel/XPU has a smaller per-thread scratch space (PTSS) limit; smaller
+        # block sizes and warp counts avoid register spilling on larger D.
+        bs_list = [512, 1024]
+        nw_list = [1, 2, 4, 8]
+    elif IS_AMD:
+        bs_list = [512, 1024, 2048, 4096, 8192]
+        nw_list = [1, 2, 4, 8, 16]
+    else:
+        bs_list = [512, 1024, 2048, 4096, 8192]
+        nw_list = [1, 2, 4, 8, 16, 32]
+    return [triton.Config({'B': bs}, num_warps=nw) for bs in bs_list for nw in nw_list]
 def _get_stride(x: torch.Tensor) -> int:
     """Get the row stride for viewing a tensor as 2D (num_rows, D) where D = shape[-1].
     Returns stride(-2) if the tensor is at least 2D, or 0 for 1D tensors.
+    The caller must ensure stride(-1) == 1, with higher dimensions contiguous relative to dimension -2.
     """
     if x.ndim < 2:
         return 0
+    if torch.compiler.is_compiling():
+        return x.shape[-1]
     return x.stride(-2)
 def _is_inner_contiguous(x: torch.Tensor) -> bool:
     """Check if a tensor can be safely viewed as 2D (num_rows, D) with row stride = stride(-2).
+    This holds when stride(-1) == 1 and all dimensions above -2 are contiguous with respect to the dimension below them.
     """
     ndim = x.ndim
     if ndim < 2:
+        return x.is_contiguous()
     if x.stride(-1) != 1:
         return False
     if ndim == 2:
 def _ensure_inner_contiguous(x: torch.Tensor) -> torch.Tensor:
     """Make the tensor inner-contiguous if it isn't already."""
+    if torch.compiler.is_compiling():
+        return x.contiguous()
+    if _is_inner_contiguous(x=x):
         return x
     return x.contiguous()
 def _alloc_output(x: torch.Tensor, contiguous: bool = False) -> torch.Tensor:
     """Allocate the output: a fresh contiguous buffer, or ``empty_like`` otherwise.
+    ``empty_like`` keeps the input's memory format only when it is dense;
+    a non-dense strided view (e.g. a ``chunk`` slice) falls back to contiguous, not the input stride.
     """
     if contiguous:
         return x.new_empty(x.shape)
     return torch.empty_like(x)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def sigmoid_fwd_kernel(x, y, stride_x_row, stride_y_row, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def sigmoid_bwd_kernel(x, dy, dx, stride_x_row, stride_dy_row, stride_dx_row, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
+@dispatch
 def sigmoid_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
+    x = _ensure_inner_contiguous(x=x)
     T, D = x.numel(), x.shape[-1]
+    y = _alloc_output(x=x, contiguous=output_contiguous)
     sigmoid_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
         T=T,
         D=D,
     )
     return y
+@dispatch
 def sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
+    x = _ensure_inner_contiguous(x=x)
+    dy = _ensure_inner_contiguous(x=dy)
     T, D = x.numel(), x.shape[-1]
+    dx = _alloc_output(x=x, contiguous=output_contiguous)
     sigmoid_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         dy=dy,
         dx=dx,
+        stride_x_row=_get_stride(x=x),
+        stride_dy_row=_get_stride(x=dy),
+        stride_dx_row=_get_stride(x=dx),
         T=T,
         D=D,
     )
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x):
         ctx.save_for_backward(x)
+        return sigmoid_fwd(x=x)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, = ctx.saved_tensors
+        return sigmoid_bwd(x=x, dy=dout)
 sigmoid = SigmoidFunction.apply
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def elu_p1_fwd_kernel(x, y, stride_x_row, stride_y_row, T, D: tl.constexpr, B: tl.constexpr):
+    i_n = tl.program_id(0).to(tl.int64)
+    offs = i_n * B + tl.arange(0, B)
+    mask = offs < T
+    row = offs // D
+    col = offs % D
+    b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32)
+    # libdevice preserves representable subnormal exponentials
+    b_y = tl.where(b_x >= 0, b_x + 1., libdevice.exp(tl.minimum(b_x, 0.)))
+    tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
+@triton.jit(do_not_specialize=['T'])
+def elu_p1_bwd_kernel(x, dy, dx, stride_x_row, stride_dy_row, stride_dx_row, T, D: tl.constexpr, B: tl.constexpr):
+    i_n = tl.program_id(0).to(tl.int64)
+    offs = i_n * B + tl.arange(0, B)
+    mask = offs < T
+    row = offs // D
+    col = offs % D
+    b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32)
+    b_dy = tl.load(dy + row * stride_dy_row + col, mask=mask, other=0.).to(tl.float32)
+    b_dx = b_dy * libdevice.exp(tl.minimum(b_x, 0.))
+    tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
+@dispatch
+def elu_p1_fwd(x: torch.Tensor) -> torch.Tensor:
+    x = x.contiguous() if x.ndim < 2 else _ensure_inner_contiguous(x=x)
+    T, D = x.numel(), x.shape[-1] if x.ndim else 1
+    y = _alloc_output(x=x)
+    if T > 0:
+        elu_p1_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
+            x=x,
+            y=y,
+            stride_x_row=_get_stride(x=x),
+            stride_y_row=_get_stride(x=y),
+            T=T,
+            D=D,
+        )
+    return y
+@dispatch
+def elu_p1_bwd(x: torch.Tensor, dy: torch.Tensor) -> torch.Tensor:
+    x = x.contiguous() if x.ndim < 2 else _ensure_inner_contiguous(x=x)
+    dy = dy.contiguous() if dy.ndim < 2 else _ensure_inner_contiguous(x=dy)
+    T, D = x.numel(), x.shape[-1] if x.ndim else 1
+    dx = _alloc_output(x=x)
+    if T > 0:
+        elu_p1_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
+            x=x,
+            dy=dy,
+            dx=dx,
+            stride_x_row=_get_stride(x=x),
+            stride_dy_row=_get_stride(x=dy),
+            stride_dx_row=_get_stride(x=dx),
+            T=T,
+            D=D,
+        )
+    return dx
+class ELUPlusOneFunction(torch.autograd.Function):
+    @staticmethod
+    @input_guard(no_guard_contiguous=True)
+    @autocast_custom_fwd
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return elu_p1_fwd(x=x)
+    @staticmethod
+    @input_guard(no_guard_contiguous=True)
+    @autocast_custom_bwd
+    def backward(ctx, dout):
+        x, = ctx.saved_tensors
+        return elu_p1_bwd(x=x, dy=dout)
+def elu_p1(x: torch.Tensor) -> torch.Tensor:
+    """Compute ELU + 1 without cancellation in the negative branch."""
+    # triton-ascend 3.2.2 does not implement libdevice.exp
+    if IS_NPU or x.device.type == 'cpu' or x.dtype == torch.float64:
+        return torch.where(x >= 0, x + 1, x.clamp_max(0).exp()).to(x.dtype)
+    return ELUPlusOneFunction.apply(x)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
+@triton.jit(do_not_specialize=['T'])
+def logsigmoid_fwd_kernel(x, y, stride_x_row, stride_y_row, temperature, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
 def logsigmoid_bwd_kernel(
     x,
     tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
+@dispatch
 def logsigmoid_fwd(x: torch.Tensor, temperature: float = 1., output_contiguous: bool = False) -> torch.Tensor:
+    x = _ensure_inner_contiguous(x=x)
     T, D = x.numel(), x.shape[-1]
+    y = _alloc_output(x=x, contiguous=output_contiguous)
     logsigmoid_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
         temperature=temperature,
         T=T,
         D=D,
     return y
+@dispatch
 def logsigmoid_bwd(
     x: torch.Tensor,
     dy: torch.Tensor,
     temperature: float = 1.,
     output_contiguous: bool = False,
 ) -> torch.Tensor:
+    x = _ensure_inner_contiguous(x=x)
+    dy = _ensure_inner_contiguous(x=dy)
     T, D = x.numel(), x.shape[-1]
+    dx = _alloc_output(x=x, contiguous=output_contiguous)
     logsigmoid_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         dy=dy,
         dx=dx,
+        stride_x_row=_get_stride(x=x),
+        stride_dy_row=_get_stride(x=dy),
+        stride_dx_row=_get_stride(x=dx),
         temperature=temperature,
         T=T,
         D=D,
     def forward(ctx, x, temperature):
         ctx.save_for_backward(x)
         ctx.temperature = temperature
+        return logsigmoid_fwd(x=x, temperature=temperature)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dy):
         x, = ctx.saved_tensors
+        return logsigmoid_bwd(x=x, dy=dy, temperature=ctx.temperature), None
 def logsigmoid(x: torch.Tensor, temperature: float = 1.) -> torch.Tensor:
     return LogSigmoidFunction.apply(x, temperature)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def swish_fwd_kernel(x, y, stride_x_row, stride_y_row, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(y + row * stride_y_row + col, b_y.to(y.dtype.element_ty), mask=mask)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def swish_bwd_kernel(x, dy, dx, stride_x_row, stride_dy_row, stride_dx_row, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask)
+@dispatch
 def swish_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
+    x = _ensure_inner_contiguous(x=x)
     T, D = x.numel(), x.shape[-1]
+    y = _alloc_output(x=x, contiguous=output_contiguous)
     swish_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
         T=T,
         D=D,
     )
     return y
+@dispatch
 def swish_bwd(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
+    x = _ensure_inner_contiguous(x=x)
+    dy = _ensure_inner_contiguous(x=dy)
     T, D = x.numel(), x.shape[-1]
+    dx = _alloc_output(x=x, contiguous=output_contiguous)
     swish_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         dy=dy,
         dx=dx,
+        stride_x_row=_get_stride(x=x),
+        stride_dy_row=_get_stride(x=dy),
+        stride_dx_row=_get_stride(x=dx),
         T=T,
         D=D,
     )
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x):
         ctx.save_for_backward(x)
+        return swish_fwd(x=x)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, = ctx.saved_tensors
+        return swish_bwd(x=x, dy=dout)
 swish = SwishFunction.apply
 # sqrt(2/pi)  -> 0.79788456
+# use the tanh GELU approximation.
 @torch.compile
 def bias_gelu(y, bias):
     x = bias + y
     return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype)
 @torch.compile
 def bias_gelu_bwd(g, y, bias):
+    """Inputs have shape [B, D] and bias has shape [D]."""
     x = bias + y
     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
     grad_y = ff * g
     return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype)
 class GeLUFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input, bias):
         ctx.save_for_backward(input, bias)
+        return bias_gelu(y=input, bias=bias)
     @staticmethod
     def backward(ctx, grad_output):
         input, bias = ctx.saved_tensors
+        return bias_gelu_bwd(g=grad_output, y=input, bias=bias)
 bias_gelu_impl = GeLUFunction.apply
+# use the tanh GELU approximation.
 @torch.compile
 def gelu_fwd(x):
     return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype)
 @torch.compile
 def gelu_bwd(g, x):
     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
     return (ff * g).to(dtype=x.dtype)
 class FastGeLUFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input):
         ctx.save_for_backward(input)
+        return gelu_fwd(x=input)
     @staticmethod
     def backward(ctx, grad_output):
         (input,) = ctx.saved_tensors
+        tmp = gelu_bwd(g=grad_output, x=input)
         return tmp
     return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype)
 @torch.compile
 def sqrelu_fwd(x):
     r = F.relu(x.float())
     return (r * r).to(dtype=x.dtype)
 @torch.compile
 def sqrelu_bwd(g, x):
     return (2.0 * g * F.relu(x.float())).to(dtype=x.dtype)
     @staticmethod
     def forward(ctx, input):
         ctx.save_for_backward(input)
+        return sqrelu_fwd(x=input)
     @staticmethod
     def backward(ctx, grad_output):
         input, = ctx.saved_tensors
+        return sqrelu_bwd(g=grad_output, x=input)
 sqrelu = SquaredReLUFunction.apply
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def swiglu_fwd_kernel(x, y, z, stride_x_row, stride_y_row, stride_z_row, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
+@triton.heuristics({'HAS_WEIGHT': lambda args: args['z'] is not None})
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
 def swiglu_fwdbwd_kernel(
+    x,
+    y,
+    g,
+    dx,
+    dy,
+    z,
     stride_x_row,
     stride_y_row,
     stride_g_row,
         tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
+@dispatch
 def swiglu_fwd(x: torch.Tensor, y: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"swiglu_fwd: shape mismatch x={x.shape} y={y.shape}"
+    x = _ensure_inner_contiguous(x=x)
+    y = _ensure_inner_contiguous(x=y)
     T, D = x.numel(), x.shape[-1]
+    z = _alloc_output(x=x, contiguous=output_contiguous)
     swiglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
         z=z,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
+        stride_z_row=_get_stride(x=z),
         T=T,
         D=D,
     )
     return z
+@dispatch
 def swiglu_fwdbwd(
     x: torch.Tensor,
     y: torch.Tensor,
     output_contiguous: bool = False,
 ):
     assert x.shape == y.shape == g.shape, f"swiglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
+    x = _ensure_inner_contiguous(x=x)
+    y = _ensure_inner_contiguous(x=y)
+    g = _ensure_inner_contiguous(x=g)
     T, D = x.numel(), x.shape[-1]
+    dx = _alloc_output(x=x, contiguous=output_contiguous)
+    dy = _alloc_output(x=y, contiguous=output_contiguous)
     if use_weight:
+        z = _alloc_output(x=x, contiguous=output_contiguous)
     else:
         z = None
     swiglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         dx=dx,
         dy=dy,
         z=z,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
+        stride_g_row=_get_stride(x=g),
+        stride_dx_row=_get_stride(x=dx),
+        stride_dy_row=_get_stride(x=dy),
+        stride_z_row=_get_stride(x=z) if z is not None else 0,
         T=T,
         D=D,
     )
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x, y):
         ctx.save_for_backward(x, y)
+        return swiglu_fwd(x=x, y=y)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, y = ctx.saved_tensors
+        return swiglu_fwdbwd(x=x, y=y, g=dout)
 class SwiGLULinearFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     @autocast_custom_fwd
     def forward(ctx, x, y, weight, bias):
+        z = swiglu_fwd(x=x, y=y, output_contiguous=True)
         out = F.linear(z, weight, bias)
         ctx.save_for_backward(x, y, weight)
         return out
     @staticmethod
     def backward(ctx, dout, *args):
         x, y, weight = ctx.saved_tensors
         dout = dout.reshape(-1, dout.shape[-1])
+        dx = dy = dlinear_weight = None
+        if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
+            dz = F.linear(dout, weight.t()).view_as(x)
+            if ctx.needs_input_grad[2]:
+                dx, dy, z = swiglu_fwdbwd(x=x, y=y, g=dz, use_weight=True, output_contiguous=True)
+            else:
+                dx, dy = swiglu_fwdbwd(x=x, y=y, g=dz, output_contiguous=True)
+        elif ctx.needs_input_grad[2]:
+            z = swiglu_fwd(x=x, y=y, output_contiguous=True)
+        if ctx.needs_input_grad[2]:
+            dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1]))
+        dlinear_bias = dout.sum(0) if ctx.needs_input_grad[3] else None
         return dx, dy, dlinear_weight, dlinear_bias
 swiglu = SwiGLUFunction.apply
+@dispatch
 def swiglu_linear(x, y, weight, bias):
     return SwiGLULinearFunction.apply(x, y, weight, bias)
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def sigmoidglu_fwd_kernel(x, y, z, stride_x_row, stride_y_row, stride_z_row, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
+@triton.heuristics({'HAS_WEIGHT': lambda args: args['z'] is not None})
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
 def sigmoidglu_fwdbwd_kernel(
+    x,
+    y,
+    g,
+    dx,
+    dy,
+    z,
     stride_x_row,
     stride_y_row,
     stride_g_row,
 @torch.compiler.disable
 def sigmoidglu_fwd(x: torch.Tensor, y: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"sigmoidglu_fwd: shape mismatch x={x.shape} y={y.shape}"
+    x = _ensure_inner_contiguous(x=x)
+    y = _ensure_inner_contiguous(x=y)
     T, D = x.numel(), x.shape[-1]
+    z = _alloc_output(x=x, contiguous=output_contiguous)
     sigmoidglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
         z=z,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
+        stride_z_row=_get_stride(x=z),
         T=T,
         D=D,
     )
     output_contiguous: bool = False,
 ):
     assert x.shape == y.shape == g.shape, f"sigmoidglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
+    x = _ensure_inner_contiguous(x=x)
+    y = _ensure_inner_contiguous(x=y)
+    g = _ensure_inner_contiguous(x=g)
     T, D = x.numel(), x.shape[-1]
+    dx = _alloc_output(x=x, contiguous=output_contiguous)
+    dy = _alloc_output(x=y, contiguous=output_contiguous)
     if use_weight:
+        z = _alloc_output(x=x, contiguous=output_contiguous)
     else:
         z = None
     sigmoidglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         dx=dx,
         dy=dy,
         z=z,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
+        stride_g_row=_get_stride(x=g),
+        stride_dx_row=_get_stride(x=dx),
+        stride_dy_row=_get_stride(x=dy),
+        stride_z_row=_get_stride(x=z) if z is not None else 0,
         T=T,
         D=D,
     )
     @input_guard(no_guard_contiguous=True)
     def forward(ctx, x, y):
         ctx.save_for_backward(x, y)
+        return sigmoidglu_fwd(x=x, y=y)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, y = ctx.saved_tensors
+        return sigmoidglu_fwdbwd(x=x, y=y, g=dout)
 class SigmoidGLULinearFunction(torch.autograd.Function):
     @input_guard(no_guard_contiguous=True)
     @autocast_custom_fwd
     def forward(ctx, x, y, weight, bias):
+        z = sigmoidglu_fwd(x=x, y=y, output_contiguous=True)
         out = F.linear(z, weight, bias)
         ctx.save_for_backward(x, y, weight)
         ctx.linear_bias_is_none = bias is None
         x, y, weight = ctx.saved_tensors
         dout = dout.reshape(-1, dout.shape[-1])
         dz = F.linear(dout, weight.t()).view_as(x)
+        dx, dy, z = sigmoidglu_fwdbwd(x=x, y=y, g=dz, use_weight=True, output_contiguous=True)
         dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1]))
         dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
         return dx, dy, dlinear_weight, dlinear_bias
 sigmoidglu_linear = SigmoidGLULinearFunction.apply
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
+def powglu_fwd_kernel(x, y, z, stride_x_row, stride_y_row, stride_z_row, m, T, D: tl.constexpr, B: tl.constexpr):
     i_n = tl.program_id(0).to(tl.int64)
     offs = i_n * B + tl.arange(0, B)
     mask = offs < T
     tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
+@triton.heuristics({'HAS_WEIGHT': lambda args: args['z'] is not None})
+@triton.autotune(configs=_activation_autotune_configs(), key=['D'], **autotune_cache_kwargs)
 @triton.jit(do_not_specialize=['T'])
 def powglu_fwdbwd_kernel(
+    x,
+    y,
+    g,
+    dx,
+    dy,
+    z,
     stride_x_row,
     stride_y_row,
     stride_g_row,
         tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask)
+@dispatch
 def powglu_fwd(x: torch.Tensor, y: torch.Tensor, power: float = 3.0, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"powglu_fwd: shape mismatch x={x.shape} y={y.shape}"
+    x = _ensure_inner_contiguous(x=x)
+    y = _ensure_inner_contiguous(x=y)
     T, D = x.numel(), x.shape[-1]
+    z = _alloc_output(x=x, contiguous=output_contiguous)
     powglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         x=x,
         y=y,
         z=z,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
+        stride_z_row=_get_stride(x=z),
         m=power,
         T=T,
         D=D,
     return z
+@dispatch
 def powglu_fwdbwd(
     x: torch.Tensor,
     y: torch.Tensor,
     output_contiguous: bool = False,
 ):
     assert x.shape == y.shape == g.shape, f"powglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
+    x = _ensure_inner_contiguous(x=x)
+    y = _ensure_inner_contiguous(x=y)
+    g = _ensure_inner_contiguous(x=g)
     T, D = x.numel(), x.shape[-1]
+    dx = _alloc_output(x=x, contiguous=output_contiguous)
+    dy = _alloc_output(x=y, contiguous=output_contiguous)
     if use_weight:
+        z = _alloc_output(x=x, contiguous=output_contiguous)
     else:
         z = None
     powglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)](
         dx=dx,
         dy=dy,
         z=z,
+        stride_x_row=_get_stride(x=x),
+        stride_y_row=_get_stride(x=y),
+        stride_g_row=_get_stride(x=g),
+        stride_dx_row=_get_stride(x=dx),
+        stride_dy_row=_get_stride(x=dy),
+        stride_z_row=_get_stride(x=z) if z is not None else 0,
         m=power,
         T=T,
         D=D,
         \text{PowGLU}(x, y) = g(x) * y,\quad
         g(x) = \begin{cases} x^{power/(\sqrt{x}+1)}\,\sigma(x) & x > 0 \\ x\,\sigma(x) & x \le 0 \end{cases}
+    For ``x <= 0`` the gate reduces to swish, matching SwiGLU; for large ``x > 0`` it saturates instead of growing,
+    replacing SwiGLU's quadratic amplification with bounded growth (Power Linear Unit, arXiv:2605.25704).
     """
     @staticmethod
     def forward(ctx, x, y, power):
         ctx.save_for_backward(x, y)
         ctx.power = power
+        return powglu_fwd(x=x, y=y, power=power)
     @staticmethod
     @input_guard(no_guard_contiguous=True)
     def backward(ctx, dout):
         x, y = ctx.saved_tensors
+        dx, dy = powglu_fwdbwd(x=x, y=y, g=dout, power=ctx.power)
         return dx, dy, None
     @input_guard(no_guard_contiguous=True)
     @autocast_custom_fwd
     def forward(ctx, x, y, weight, bias, power):
+        z = powglu_fwd(x=x, y=y, power=power, output_contiguous=True)
         out = F.linear(z, weight, bias)
         ctx.save_for_backward(x, y, weight)
         ctx.linear_bias_is_none = bias is None
         x, y, weight = ctx.saved_tensors
         dout = dout.reshape(-1, dout.shape[-1])
         dz = F.linear(dout, weight.t()).view_as(x)
+        dx, dy, z = powglu_fwdbwd(x=x, y=y, g=dz, power=ctx.power, use_weight=True, output_contiguous=True)
         dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1]))
         dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
         return dx, dy, dlinear_weight, dlinear_bias, None
     return PowGLUFunction.apply(x, y, power)
+@dispatch
 def powglu_linear(
     x: torch.Tensor,
     y: torch.Tensor,

build/torch-cuda/modules/backends/__init__.py DELETED Viewed

@@ -1,17 +0,0 @@
-# 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
-"""Module-level backends for FLA components such as rotary and cross-entropy."""
-from ...modules.backends.triton_ascend import TritonAscendBackend
-from ...ops.backends import BackendRegistry, dispatch
-modules_registry = BackendRegistry("modules")
-modules_registry.register(TritonAscendBackend())
-__all__ = ['dispatch', 'modules_registry']

build/torch-cuda/modules/backends/triton_ascend/__init__.py DELETED Viewed

@@ -1,422 +0,0 @@
-# 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
-"""Triton-Ascend (Huawei NPU) backend for FLA modules."""
-from __future__ import annotations
-from ....ops.backends import BaseBackend
-class TritonAscendBackend(BaseBackend):
-    """Ascend NPU backend using triton-ascend kernels."""
-    backend_type = "triton_ascend"
-    package_name = None
-    env_var = None
-    priority = 0
-    @classmethod
-    def is_available(cls) -> bool:
-        from ....utils import IS_NPU
-        return IS_NPU
-    def rotary_embedding_fwdbwd(
-        self,
-        x,
-        cos,
-        sin,
-        seqlen_offsets=0,
-        cu_seqlens=None,
-        interleaved=False,
-        inplace=False,
-        conjugate=False,
-        chunk_indices=None,
-    ):
-        from ....modules.backends.triton_ascend.rotary import rotary_embedding_fwdbwd_npu
-        return rotary_embedding_fwdbwd_npu(
-            x,
-            cos,
-            sin,
-            seqlen_offsets=seqlen_offsets,
-            cu_seqlens=cu_seqlens,
-            interleaved=interleaved,
-            inplace=inplace,
-            conjugate=conjugate,
-            chunk_indices=chunk_indices,
-        )
-    def cross_entropy_loss(
-        self,
-        logits,
-        target,
-        label_smoothing=0.0,
-        logit_scale=1.0,
-        lse_square_scale=0.0,
-        logit_softcapping=None,
-        ignore_index=-100,
-        inplace_backward=False,
-        process_group=None,
-    ):
-        from ....modules.backends.triton_ascend.fused_cross_entropy import (
-            cross_entropy_loss_npu,
-        )
-        return cross_entropy_loss_npu(
-            logits,
-            target,
-            label_smoothing,
-            logit_scale,
-            lse_square_scale,
-            logit_softcapping,
-            ignore_index,
-            inplace_backward,
-            process_group,
-        )
-    def logsumexp_fwd(
-        self,
-        x,
-        scale=None,
-        softcapping=None,
-        dtype=None,
-    ):
-        from ....modules.backends.triton_ascend.fused_linear_cross_entropy import (
-            logsumexp_fwd_npu,
-        )
-        return logsumexp_fwd_npu(x, scale=scale, softcapping=softcapping, dtype=dtype)
-    def fused_linear_cross_entropy_forward(
-        self,
-        x,
-        target,
-        weight,
-        bias=None,
-        ignore_index=-100,
-        label_smoothing=0.0,
-        logit_scale=1.0,
-        logit_softcapping=None,
-        num_chunks=8,
-        reduction="mean",
-        use_l2warp=False,
-        l2_penalty_factor=1e-4,
-        accumulate_grad_in_fp32=True,
-    ):
-        from ....modules.backends.triton_ascend.fused_linear_cross_entropy import (
-            fused_linear_cross_entropy_forward_npu,
-        )
-        return fused_linear_cross_entropy_forward_npu(
-            x,
-            target,
-            weight,
-            bias,
-            ignore_index,
-            label_smoothing,
-            logit_scale,
-            logit_softcapping,
-            num_chunks,
-            reduction,
-            use_l2warp,
-            l2_penalty_factor,
-            accumulate_grad_in_fp32,
-        )
-    def fused_linear_cross_entropy_backward(
-        self,
-        do,
-        dx,
-        dw,
-        db,
-    ):
-        from ....modules.backends.triton_ascend.fused_linear_cross_entropy import (
-            fused_linear_cross_entropy_backward_npu,
-        )
-        return fused_linear_cross_entropy_backward_npu(do, dx, dw, db)
-    def sigmoid_fwd(self, x, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import sigmoid_fwd_npu
-        return sigmoid_fwd_npu(x, output_contiguous=output_contiguous)
-    def sigmoid_bwd(self, x, dy, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import sigmoid_bwd_npu
-        return sigmoid_bwd_npu(x, dy, output_contiguous=output_contiguous)
-    def logsigmoid_fwd(self, x, temperature=1., output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import logsigmoid_fwd_npu
-        return logsigmoid_fwd_npu(x, temperature=temperature, output_contiguous=output_contiguous)
-    def logsigmoid_bwd(self, x, dy, temperature=1., output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import logsigmoid_bwd_npu
-        return logsigmoid_bwd_npu(x, dy, temperature=temperature, output_contiguous=output_contiguous)
-    def swish_fwd(self, x, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import swish_fwd_npu
-        return swish_fwd_npu(x, output_contiguous=output_contiguous)
-    def swish_bwd(self, x, dy, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import swish_bwd_npu
-        return swish_bwd_npu(x, dy, output_contiguous=output_contiguous)
-    def swiglu_fwd(self, x, y, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import swiglu_fwd_npu
-        return swiglu_fwd_npu(x, y, output_contiguous=output_contiguous)
-    def swiglu_fwdbwd(self, x, y, g, use_weight=False, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import swiglu_fwdbwd_npu
-        return swiglu_fwdbwd_npu(x, y, g, use_weight=use_weight, output_contiguous=output_contiguous)
-    def swiglu_linear(self, x, y, weight, bias):
-        from ....modules.backends.triton_ascend.activations import swiglu_linear_npu
-        return swiglu_linear_npu(x, y, weight, bias)
-    def gelu_fwd(self, x):
-        from ....modules.backends.triton_ascend.activations import gelu_fwd_npu
-        return gelu_fwd_npu(x)
-    def gelu_bwd(self, g, x):
-        from ....modules.backends.triton_ascend.activations import gelu_bwd_npu
-        return gelu_bwd_npu(g, x)
-    def sqrelu_fwd(self, x):
-        from ....modules.backends.triton_ascend.activations import sqrelu_fwd_npu
-        return sqrelu_fwd_npu(x)
-    def sqrelu_bwd(self, g, x):
-        from ....modules.backends.triton_ascend.activations import sqrelu_bwd_npu
-        return sqrelu_bwd_npu(g, x)
-    def powglu_fwd(self, x, y, power=3.0, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import powglu_fwd_npu
-        return powglu_fwd_npu(x, y, power=power, output_contiguous=output_contiguous)
-    def powglu_fwdbwd(self, x, y, g, power=3.0, use_weight=False, output_contiguous=False):
-        from ....modules.backends.triton_ascend.activations import powglu_fwdbwd_npu
-        return powglu_fwdbwd_npu(
-            x, y, g, power=power, use_weight=use_weight, output_contiguous=output_contiguous,
-        )
-    def powglu_linear(self, x, y, weight, bias, power=3.0):
-        from ....modules.backends.triton_ascend.activations import powglu_linear_npu
-        return powglu_linear_npu(x, y, weight, bias, power)
-    def fused_kl_div_forward(
-        self,
-        x,
-        target_x,
-        weight,
-        target_weight,
-        reduction='batchmean',
-        accumulate_grad_in_fp32=True,
-    ):
-        from ....modules.backends.triton_ascend.fused_kl_div import fused_kl_div_forward_npu
-        return fused_kl_div_forward_npu(
-            x,
-            target_x,
-            weight,
-            target_weight,
-            reduction,
-            accumulate_grad_in_fp32,
-        )
-    def fused_kl_div_backward(self, do, dx, dw):
-        from ....modules.backends.triton_ascend.fused_kl_div import fused_kl_div_backward_npu
-        return fused_kl_div_backward_npu(do, dx, dw)
-    def layer_norm_fwd(
-        self,
-        x,
-        weight,
-        bias,
-        eps=1e-5,
-        residual=None,
-        out_dtype=None,
-        residual_dtype=None,
-        is_rms_norm=False,
-        num_groups=1,
-    ):
-        from ....modules.backends.triton_ascend.layernorm import layer_norm_fwd_npu
-        return layer_norm_fwd_npu(
-            x,
-            weight,
-            bias,
-            eps,
-            residual,
-            out_dtype,
-            residual_dtype,
-            is_rms_norm,
-            num_groups,
-        )
-    def layer_norm_bwd(
-        self,
-        dy,
-        x,
-        weight,
-        bias,
-        mean=None,
-        rstd=None,
-        dres=None,
-        has_residual=False,
-        is_rms_norm=False,
-        x_dtype=None,
-        recompute_output=False,
-        num_groups=1,
-    ):
-        from ....modules.backends.triton_ascend.layernorm import layer_norm_bwd_npu
-        return layer_norm_bwd_npu(
-            dy,
-            x,
-            weight,
-            bias,
-            mean,
-            rstd,
-            dres,
-            has_residual,
-            is_rms_norm,
-            x_dtype,
-            recompute_output,
-            num_groups,
-        )
-    def fused_grpo_loss(
-        self,
-        logits,
-        ref_logp,
-        input_ids,
-        advantages,
-        beta=0.1,
-        completion_mask=None,
-        save_kl=False,
-        inplace=False,
-    ):
-        from ....modules.backends.triton_ascend.grpo import fused_grpo_loss_npu
-        return fused_grpo_loss_npu(
-            logits,
-            ref_logp,
-            input_ids,
-            advantages,
-            beta,
-            completion_mask,
-            save_kl,
-            inplace,
-        )
-    def causal_conv1d_fwd(
-        self,
-        x,
-        weight,
-        bias,
-        residual,
-        initial_state=None,
-        output_final_state=False,
-        activation=None,
-        cu_seqlens=None,
-        cu_seqlens_cpu=None,
-        chunk_indices=None,
-        BT=64,
-        layout_fallback=False,
-    ):
-        from ....modules.backends.triton_ascend.causal_conv1d import causal_conv1d_fwd_npu
-        return causal_conv1d_fwd_npu(
-            x,
-            weight,
-            bias,
-            residual,
-            initial_state,
-            output_final_state,
-            activation,
-            cu_seqlens,
-            cu_seqlens_cpu,
-            chunk_indices,
-            BT,
-            layout_fallback,
-        )
-    def causal_conv1d_bwd(
-        self,
-        x,
-        dy,
-        dht,
-        weight=None,
-        bias=None,
-        residual=None,
-        initial_state=None,
-        activation=None,
-        cu_seqlens=None,
-        cu_seqlens_cpu=None,
-        chunk_indices=None,
-        BT=64,
-        layout_fallback=False,
-    ):
-        from ....modules.backends.triton_ascend.causal_conv1d import causal_conv1d_bwd_npu
-        return causal_conv1d_bwd_npu(
-            x,
-            dy,
-            dht,
-            weight,
-            bias,
-            residual,
-            initial_state,
-            activation,
-            cu_seqlens,
-            cu_seqlens_cpu,
-            chunk_indices,
-            BT,
-            layout_fallback,
-        )
-    def compute_dh0_triton(
-        self,
-        dy,
-        y,
-        weight,
-        initial_state,
-        activation,
-        cu_seqlens,
-    ):
-        from ....modules.backends.triton_ascend.causal_conv1d import compute_dh0_npu
-        return compute_dh0_npu(
-            dy,
-            y,
-            weight,
-            initial_state,
-            activation,
-            cu_seqlens,
-        )
-    def causal_conv1d_update_states(
-        self,
-        x,
-        state_len,
-        initial_state=None,
-        cu_seqlens=None,
-    ):
-        from ....modules.backends.triton_ascend.causal_conv1d import causal_conv1d_update_states_npu
-        return causal_conv1d_update_states_npu(
-            x,
-            state_len,
-            initial_state,
-            cu_seqlens,
-        )
-    def causal_conv1d_update(
-        self,
-        x,
-        cache,
-        residual=None,
-        weight=None,
-        bias=None,
-        activation=None,
-    ):
-        from ....modules.backends.triton_ascend.causal_conv1d import causal_conv1d_update_npu
-        return causal_conv1d_update_npu(
-            x,
-            cache,
-            residual,
-            weight,
-            bias,
-            activation,
-        )

build/torch-cuda/modules/backends/triton_ascend/activations.py DELETED Viewed

@@ -1,931 +0,0 @@
-# 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
-"""Activation kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import torch.nn.functional as F
-import triton
-import triton.language as tl
-from ....ops.utils.op import exp, log
-from ....utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
-from ....utils.ascend_ub_manager import ASCEND_MAX_GRID_DIM, compute_activation_block_size
-# Ascend launch limits: grid dim and per-core vector width.
-_MAX_CORE_DIM = 65535
-def _activation_launch_config(
-    T: int,
-    is_backward: bool = False,
-    *,
-    memory_multiplier: float | None = None,
-) -> tuple[tuple[int], int]:
-    """Pick block size under Ascend launch and UB limits."""
-    B = compute_activation_block_size(
-        T,
-        is_backward,
-        max_grid=ASCEND_MAX_GRID_DIM,
-        max_core_dim=_MAX_CORE_DIM,
-        memory_multiplier=memory_multiplier,
-    )
-    return (triton.cdiv(T, B),), B
-@triton.jit
-def _flat_offset(
-    offs,
-    D: tl.constexpr,
-    stride,
-    IS_LINEAR: tl.constexpr,
-):
-    if IS_LINEAR:
-        return offs
-    row = offs // D
-    col = offs % D
-    return row * stride + col
-def _get_stride(x: torch.Tensor) -> int:
-    if x.ndim < 2:
-        return 0
-    return x.stride(-2)
-def _is_linear_stride(stride: int, D: int) -> bool:
-    return stride == D
-_LINEAR_HEURISTICS_XY = {
-    'X_LINEAR': lambda args: _is_linear_stride(args['stride_x_row'], args['D']),
-    'Y_LINEAR': lambda args: _is_linear_stride(args['stride_y_row'], args['D']),
-}
-_LINEAR_HEURISTICS_XYZ = {
-    **_LINEAR_HEURISTICS_XY,
-    'Z_LINEAR': lambda args: _is_linear_stride(args['stride_z_row'], args['D']),
-}
-_LINEAR_HEURISTICS_BWD = {
-    'X_LINEAR': lambda args: _is_linear_stride(args['stride_x_row'], args['D']),
-    'DY_LINEAR': lambda args: _is_linear_stride(args['stride_dy_row'], args['D']),
-    'DX_LINEAR': lambda args: _is_linear_stride(args['stride_dx_row'], args['D']),
-}
-_LINEAR_HEURISTICS_FWDBWD = {
-    **_LINEAR_HEURISTICS_XYZ,
-    'G_LINEAR': lambda args: _is_linear_stride(args['stride_g_row'], args['D']),
-    'DX_LINEAR': lambda args: _is_linear_stride(args['stride_dx_row'], args['D']),
-    'DY_LINEAR': lambda args: _is_linear_stride(args['stride_dy_row'], args['D']),
-}
-def _is_inner_contiguous(x: torch.Tensor) -> bool:
-    ndim = x.ndim
-    if ndim < 2:
-        return True
-    if x.stride(-1) != 1:
-        return False
-    if ndim == 2:
-        return True
-    if ndim == 3:
-        return x.stride(0) == x.stride(-2) * x.shape[-2]
-    if ndim == 4:
-        if x.stride(1) != x.stride(-2) * x.shape[-2]:
-            return False
-        return x.stride(0) == x.stride(1) * x.shape[1]
-    expected = x.stride(-2) * x.shape[-2]
-    for d in range(ndim - 3, -1, -1):
-        if x.stride(d) != expected:
-            return False
-        expected *= x.shape[d]
-    return True
-def _ensure_inner_contiguous(x: torch.Tensor) -> torch.Tensor:
-    if _is_inner_contiguous(x):
-        return x
-    return x.contiguous()
-def _alloc_output(x: torch.Tensor, contiguous: bool = False) -> torch.Tensor:
-    if contiguous:
-        return x.new_empty(x.shape)
-    return torch.empty_like(x)
-@triton.heuristics(_LINEAR_HEURISTICS_XY)
-@triton.jit(do_not_specialize=['T'])
-def sigmoid_fwd_kernel(
-    x, y,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    y_val = tl.sigmoid(x_val)
-    tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_BWD)
-@triton.jit(do_not_specialize=['T'])
-def sigmoid_bwd_kernel(
-    x, dy, dx,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_dy_row,
-    stride_dx_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
-    s = tl.sigmoid(x_val)
-    dx_val = g_val * s * (1.0 - s)
-    tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_XY)
-@triton.jit(do_not_specialize=['T'])
-def logsigmoid_fwd_kernel(
-    x,
-    y,
-    temperature,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-):
-    i = tl.program_id(0)
-    offs = i * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    b_x = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    b_m = tl.minimum(0., b_x)
-    b_z = 1. + exp(-tl.abs(b_x))
-    b_y = (b_m - log(b_z)) / temperature
-    tl.store(y + y_off, b_y.to(y.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_BWD)
-@triton.jit(do_not_specialize=['T'])
-def logsigmoid_bwd_kernel(
-    x,
-    dx,
-    dy,
-    temperature,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_dx_row,
-    stride_dy_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-):
-    i = tl.program_id(0)
-    offs = i * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    b_x = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    b_dy = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
-    b_s = tl.sigmoid(b_x)
-    b_dx = b_dy * ((1. - b_s) / temperature)
-    tl.store(dx + dx_off, b_dx.to(dx.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_XY)
-@triton.jit(do_not_specialize=['T'])
-def swish_fwd_kernel(
-    x, y,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    s = tl.sigmoid(x_val)
-    y_val = x_val * s
-    tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_BWD)
-@triton.jit(do_not_specialize=['T'])
-def swish_bwd_kernel(
-    x, dy, dx,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_dy_row,
-    stride_dx_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
-    s = tl.sigmoid(x_val)
-    dx_val = g_val * s * (1.0 + x_val * (1.0 - s))
-    tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_XYZ)
-@triton.jit(do_not_specialize=['T'])
-def swiglu_fwd_kernel(
-    x, y, z,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    stride_z_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-    Z_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    z_off = _flat_offset(offs, D, stride_z_row, Z_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    y_val = tl.load(y + y_off, mask=mask, other=0.).to(tl.float32)
-    s = tl.sigmoid(x_val)
-    z_val = x_val * s * y_val
-    tl.store(z + z_off, z_val.to(z.dtype.element_ty), mask=mask)
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['z'] is not None,
-    **_LINEAR_HEURISTICS_FWDBWD,
-})
-@triton.jit(do_not_specialize=['T'])
-def swiglu_fwdbwd_kernel(
-    x, y, g, dx, dy, z,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    stride_g_row,
-    stride_dx_row,
-    stride_dy_row,
-    stride_z_row,
-    B: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-    G_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-    Z_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    g_off = _flat_offset(offs, D, stride_g_row, G_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    y_val = tl.load(y + y_off, mask=mask, other=0.).to(tl.float32)
-    g_val = tl.load(g + g_off, mask=mask, other=0.).to(tl.float32)
-    s = tl.sigmoid(x_val)
-    x_s = x_val * s
-    dx_val = g_val * s * (1.0 + x_val * (1.0 - s)) * y_val
-    dy_val = g_val * x_s
-    tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)
-    tl.store(dy + dy_off, dy_val.to(dy.dtype.element_ty), mask=mask)
-    if HAS_WEIGHT:
-        z_off = _flat_offset(offs, D, stride_z_row, Z_LINEAR)
-        z_val = x_s * y_val
-        tl.store(z + z_off, z_val.to(z.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_XY)
-@triton.jit(do_not_specialize=['T'])
-def gelu_fwd_kernel(
-    x, y,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    t = 0.79788456 * x_val * (1.0 + 0.044715 * x_val * x_val)
-    tanh_out = tl.tanh(t)
-    y_val = x_val * 0.5 * (1.0 + tanh_out)
-    tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_BWD)
-@triton.jit(do_not_specialize=['T'])
-def gelu_bwd_kernel(
-    x, dy, dx,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_dy_row,
-    stride_dx_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
-    t = 0.79788456 * x_val * (1.0 + 0.044715 * x_val * x_val)
-    tanh_out = tl.tanh(t)
-    ff = 0.5 * x_val * (
-        (1.0 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x_val * x_val)
-    ) + 0.5 * (1.0 + tanh_out)
-    dx_val = ff * g_val
-    tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_XY)
-@triton.jit(do_not_specialize=['T'])
-def sqrelu_fwd_kernel(
-    x, y,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_y_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    r = tl.maximum(x_val, 0.0)
-    y_val = r * r
-    tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask)
-@triton.heuristics(_LINEAR_HEURISTICS_BWD)
-@triton.jit(do_not_specialize=['T'])
-def sqrelu_bwd_kernel(
-    x, dy, dx,
-    T,
-    D: tl.constexpr,
-    stride_x_row,
-    stride_dy_row,
-    stride_dx_row,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
-    dx_val = 2.0 * g_val * tl.maximum(x_val, 0.0)
-    tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)
-@torch.compiler.disable
-def gelu_fwd_npu(x: torch.Tensor) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x)
-    grid, B = _activation_launch_config(T)
-    gelu_fwd_kernel[grid](
-        x, y, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        BLOCK_SIZE=B,
-    )
-    return y
-@torch.compiler.disable
-def gelu_bwd_npu(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    g = _ensure_inner_contiguous(g)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x)
-    grid, B = _activation_launch_config(T, is_backward=True)
-    gelu_bwd_kernel[grid](
-        x, g, dx, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        BLOCK_SIZE=B,
-    )
-    return dx
-@torch.compiler.disable
-def sqrelu_fwd_npu(x: torch.Tensor) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x)
-    grid, B = _activation_launch_config(T)
-    sqrelu_fwd_kernel[grid](
-        x, y, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        BLOCK_SIZE=B,
-    )
-    return y
-@torch.compiler.disable
-def sqrelu_bwd_npu(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    g = _ensure_inner_contiguous(g)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x)
-    grid, B = _activation_launch_config(T, is_backward=True)
-    sqrelu_bwd_kernel[grid](
-        x, g, dx, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        BLOCK_SIZE=B,
-    )
-    return dx
-@torch.compiler.disable
-def sigmoid_fwd_npu(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
-    sigmoid_fwd_kernel[grid](
-        x, y, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        B=B,
-    )
-    return y
-@torch.compiler.disable
-def sigmoid_bwd_npu(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    dy = _ensure_inner_contiguous(dy)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T, is_backward=True)
-    sigmoid_bwd_kernel[grid](
-        x, dy, dx, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(dy),
-        stride_dx_row=_get_stride(dx),
-        B=B,
-    )
-    return dx
-@torch.compiler.disable
-def logsigmoid_fwd_npu(x: torch.Tensor, temperature: float = 1., output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
-    logsigmoid_fwd_kernel[grid](
-        x=x,
-        y=y,
-        temperature=temperature,
-        T=T,
-        D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        B=B,
-    )
-    return y
-@torch.compiler.disable
-def logsigmoid_bwd_npu(
-    x: torch.Tensor,
-    dy: torch.Tensor,
-    temperature: float = 1.,
-    output_contiguous: bool = False,
-) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    dy = _ensure_inner_contiguous(dy)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T, is_backward=True)
-    logsigmoid_bwd_kernel[grid](
-        x=x,
-        dx=dx,
-        dy=dy,
-        temperature=temperature,
-        T=T,
-        D=D,
-        stride_x_row=_get_stride(x),
-        stride_dx_row=_get_stride(dx),
-        stride_dy_row=_get_stride(dy),
-        B=B,
-    )
-    return dx
-@torch.compiler.disable
-def swish_fwd_npu(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    T, D = x.numel(), x.shape[-1]
-    y = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
-    swish_fwd_kernel[grid](
-        x, y, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        B=B,
-    )
-    return y
-@torch.compiler.disable
-def swish_bwd_npu(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    x = _ensure_inner_contiguous(x)
-    dy = _ensure_inner_contiguous(dy)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T, is_backward=True)
-    swish_bwd_kernel[grid](
-        x, dy, dx, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_dy_row=_get_stride(dy),
-        stride_dx_row=_get_stride(dx),
-        B=B,
-    )
-    return dx
-@torch.compiler.disable
-def swiglu_fwd_npu(x: torch.Tensor, y: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
-    assert x.shape == y.shape, f"swiglu_fwd: shape mismatch x={x.shape} y={y.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    T, D = x.numel(), x.shape[-1]
-    z = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
-    swiglu_fwd_kernel[grid](
-        x, y, z, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_z_row=_get_stride(z),
-        B=B,
-    )
-    return z
-@torch.compiler.disable
-def swiglu_fwdbwd_npu(
-    x: torch.Tensor,
-    y: torch.Tensor,
-    g: torch.Tensor,
-    use_weight: bool = False,
-    output_contiguous: bool = False,
-):
-    assert x.shape == y.shape == g.shape, f"swiglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    g = _ensure_inner_contiguous(g)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    dy = _alloc_output(y, output_contiguous)
-    if use_weight:
-        z = _alloc_output(x, output_contiguous)
-    else:
-        z = None
-    grid, B = _activation_launch_config(T, is_backward=True)
-    swiglu_fwdbwd_kernel[grid](
-        x, y, g, dx, dy, z, T=T, D=D,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_g_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        stride_dy_row=_get_stride(dy),
-        stride_z_row=_get_stride(z) if z is not None else 0,
-        B=B,
-    )
-    if use_weight:
-        return dx, dy, z
-    return dx, dy
-class SwiGLULinearFunctionNPU(torch.autograd.Function):
-    @staticmethod
-    @input_guard(no_guard_contiguous=True)
-    @autocast_custom_fwd
-    def forward(ctx, x, y, weight, bias):
-        z = swiglu_fwd_npu(x, y, output_contiguous=True)
-        out = F.linear(z, weight, bias)
-        ctx.save_for_backward(x, y, weight)
-        ctx.linear_bias_is_none = bias is None
-        return out
-    @staticmethod
-    @input_guard(no_guard_contiguous=True)
-    @autocast_custom_bwd
-    def backward(ctx, dout, *args):
-        x, y, weight = ctx.saved_tensors
-        dout = dout.reshape(-1, dout.shape[-1])
-        dz = F.linear(dout, weight.t()).view_as(x)
-        dx, dy, z = swiglu_fwdbwd_npu(x, y, dz, use_weight=True, output_contiguous=True)
-        z_flat = z.reshape(-1, z.shape[-1])
-        dlinear_weight = dout.t() @ z_flat
-        dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
-        return dx, dy, dlinear_weight, dlinear_bias
-def swiglu_linear_npu(x, y, weight, bias):
-    return SwiGLULinearFunctionNPU.apply(x, y, weight, bias)
-@triton.heuristics(_LINEAR_HEURISTICS_XYZ)
-@triton.jit(do_not_specialize=['T'])
-def powglu_fwd_kernel(
-    x, y, z,
-    stride_x_row,
-    stride_y_row,
-    stride_z_row,
-    m,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-    Z_LINEAR: tl.constexpr,
-):
-    i_n = tl.program_id(0)
-    offs = i_n * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    z_off = _flat_offset(offs, D, stride_z_row, Z_LINEAR)
-    b_x = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    b_y = tl.load(y + y_off, mask=mask, other=0.).to(tl.float32)
-    b_s = tl.sigmoid(b_x)
-    b_pos = b_x > 0
-    # feed only positive lanes to log/sqrt; masked lanes give x**p = 1 and are dropped by the where
-    b_xp = tl.where(b_pos, b_x, 1.0)
-    b_sqrt = tl.sqrt(b_xp)
-    b_p = m / (b_sqrt + 1.0)
-    b_pow = exp(b_p * log(b_xp))
-    b_g = tl.where(b_pos, b_pow * b_s, b_x * b_s)
-    b_z = b_g * b_y
-    tl.store(z + z_off, b_z.to(z.dtype.element_ty), mask=mask)
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['z'] is not None,
-    **_LINEAR_HEURISTICS_FWDBWD,
-})
-@triton.jit(do_not_specialize=['T'])
-def powglu_fwdbwd_kernel(
-    x, y, g, dx, dy, z,
-    stride_x_row,
-    stride_y_row,
-    stride_g_row,
-    stride_dx_row,
-    stride_dy_row,
-    stride_z_row,
-    m,
-    T,
-    D: tl.constexpr,
-    B: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    X_LINEAR: tl.constexpr,
-    Y_LINEAR: tl.constexpr,
-    G_LINEAR: tl.constexpr,
-    DX_LINEAR: tl.constexpr,
-    DY_LINEAR: tl.constexpr,
-    Z_LINEAR: tl.constexpr,
-):
-    i_n = tl.program_id(0)
-    offs = i_n * B + tl.arange(0, B)
-    mask = offs < T
-    x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
-    y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
-    g_off = _flat_offset(offs, D, stride_g_row, G_LINEAR)
-    dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
-    dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
-    b_x = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-    b_y = tl.load(y + y_off, mask=mask, other=0.).to(tl.float32)
-    b_g = tl.load(g + g_off, mask=mask, other=0.).to(tl.float32)
-    b_s = tl.sigmoid(b_x)
-    b_pos = b_x > 0
-    b_xp = tl.where(b_pos, b_x, 1.0)
-    b_sqrt = tl.sqrt(b_xp)
-    b_ln = log(b_xp)
-    b_p = m / (b_sqrt + 1.0)
-    b_pow = exp(b_p * b_ln)
-    b_gate_pos = b_pow * b_s
-    # d/dx of the exponent term: p' = -m / (2*sqrt(x)*(sqrt(x)+1)**2)
-    b_pprime = -m / (2.0 * b_sqrt * (b_sqrt + 1.0) * (b_sqrt + 1.0))
-    b_dgate_pos = b_gate_pos * (b_pprime * b_ln + b_p / b_xp + 1.0 - b_s)
-    b_gate_neg = b_x * b_s
-    b_dgate_neg = b_s * (1.0 + b_x * (1.0 - b_s))
-    b_gate = tl.where(b_pos, b_gate_pos, b_gate_neg)
-    b_dgate = tl.where(b_pos, b_dgate_pos, b_dgate_neg)
-    b_dx = b_g * b_y * b_dgate
-    b_dy = b_g * b_gate
-    tl.store(dx + dx_off, b_dx.to(dx.dtype.element_ty), mask=mask)
-    tl.store(dy + dy_off, b_dy.to(dy.dtype.element_ty), mask=mask)
-    if HAS_WEIGHT:
-        b_z = b_gate * b_y
-        z_off = _flat_offset(offs, D, stride_z_row, Z_LINEAR)
-        tl.store(z + z_off, b_z.to(z.dtype.element_ty), mask=mask)
-# Peak fp32 temporaries: sigmoid, sqrt, log, exp, pow, gate, output.
-_POWGLU_FWD_MEM_MULT = 8.0
-_POWGLU_BWD_MEM_MULT = 10.0
-@torch.compiler.disable
-def powglu_fwd_npu(x: torch.Tensor, y: torch.Tensor, power: float = 3.0, output_contiguous: bool = False) -> torch.Tensor:
-    assert x.shape == y.shape, f"powglu_fwd: shape mismatch x={x.shape} y={y.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    T, D = x.numel(), x.shape[-1]
-    z = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T, memory_multiplier=_POWGLU_FWD_MEM_MULT)
-    powglu_fwd_kernel[grid](
-        x=x,
-        y=y,
-        z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_z_row=_get_stride(z),
-        m=power,
-        T=T,
-        D=D,
-        B=B,
-    )
-    return z
-@torch.compiler.disable
-def powglu_fwdbwd_npu(
-    x: torch.Tensor,
-    y: torch.Tensor,
-    g: torch.Tensor,
-    power: float = 3.0,
-    use_weight: bool = False,
-    output_contiguous: bool = False,
-):
-    assert x.shape == y.shape == g.shape, f"powglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}"
-    x = _ensure_inner_contiguous(x)
-    y = _ensure_inner_contiguous(y)
-    g = _ensure_inner_contiguous(g)
-    T, D = x.numel(), x.shape[-1]
-    dx = _alloc_output(x, output_contiguous)
-    dy = _alloc_output(y, output_contiguous)
-    if use_weight:
-        z = _alloc_output(x, output_contiguous)
-    else:
-        z = None
-    grid, B = _activation_launch_config(T, is_backward=True, memory_multiplier=_POWGLU_BWD_MEM_MULT)
-    powglu_fwdbwd_kernel[grid](
-        x=x,
-        y=y,
-        g=g,
-        dx=dx,
-        dy=dy,
-        z=z,
-        stride_x_row=_get_stride(x),
-        stride_y_row=_get_stride(y),
-        stride_g_row=_get_stride(g),
-        stride_dx_row=_get_stride(dx),
-        stride_dy_row=_get_stride(dy),
-        stride_z_row=_get_stride(z) if z is not None else 0,
-        m=power,
-        T=T,
-        D=D,
-        B=B,
-    )
-    if use_weight:
-        return dx, dy, z
-    return dx, dy
-class PowGLULinearFunctionNPU(torch.autograd.Function):
-    r"""
-    Power-Gated Linear Unit (PowGLU) function followed by a linear transformation.
-    .. math::
-        \text{PowGLULinear}(x, y, W, b) = (g(x) * y) W + b
-    This simple wrap discards the intermediate results of PowGLU(x, y) to save memory.
-    """
-    @staticmethod
-    @input_guard(no_guard_contiguous=True)
-    @autocast_custom_fwd
-    def forward(ctx, x, y, weight, bias, power):
-        z = powglu_fwd_npu(x, y, power, output_contiguous=True)
-        out = F.linear(z, weight, bias)
-        ctx.save_for_backward(x, y, weight)
-        ctx.linear_bias_is_none = bias is None
-        ctx.power = power
-        return out
-    @staticmethod
-    @input_guard(no_guard_contiguous=True)
-    @autocast_custom_bwd
-    def backward(ctx, dout, *args):
-        x, y, weight = ctx.saved_tensors
-        dout = dout.reshape(-1, dout.shape[-1])
-        dz = F.linear(dout, weight.t()).view_as(x)
-        dx, dy, z = powglu_fwdbwd_npu(x, y, dz, ctx.power, use_weight=True, output_contiguous=True)
-        z_flat = z.reshape(-1, z.shape[-1])
-        dlinear_weight = dout.t() @ z_flat
-        dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
-        return dx, dy, dlinear_weight, dlinear_bias, None
-def powglu_linear_npu(
-    x: torch.Tensor,
-    y: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    power: float = 3.0,
-) -> torch.Tensor:
-    return PowGLULinearFunctionNPU.apply(x, y, weight, bias, power)

build/torch-cuda/modules/backends/triton_ascend/causal_conv1d.py DELETED Viewed

@@ -1,1175 +0,0 @@
-# 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
-"""Causal 1D convolution kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import triton
-import triton.language as tl
-from einops import rearrange
-from ....ops.utils import prepare_chunk_indices
-from ....utils import input_guard
-STATIC_WARPS = 2
-# Ascend Triton rejects grids whose product exceeds 65535 (see fla/modules/token_shift.py).
-_NPU_MAX_TRITON_GRID = 65535
-_ELEM_BLOCK = 2048
-def _elementwise_launch_iters(numel: int):
-    n_blocks = triton.cdiv(numel, _ELEM_BLOCK)
-    for block_off in range(0, n_blocks, _NPU_MAX_TRITON_GRID):
-        yield min(_NPU_MAX_TRITON_GRID, n_blocks - block_off), block_off * _ELEM_BLOCK
-def _npu_chunk_size(T: int, BT: int) -> int:
-    BT = min(max(BT, 1), 64)
-    if BT not in (1, 2, 4, 8, 16, 32, 64):
-        BT = triton.next_power_of_2(BT)
-    # Ascend compiler requires power-of-2 BT; pad with mask when BT > T.
-    if T not in (1, 2, 4, 8, 16, 32, 64):
-        BT = min(triton.next_power_of_2(T), 64)
-    else:
-        BT = min(BT, T, 64)
-    return BT
-def _clamp_bd_for_grid(B: int, NT: int, D: int, BD: int) -> int:
-    while triton.cdiv(D, BD) * NT * B > _NPU_MAX_TRITON_GRID and BD < 64:
-        BD *= 2
-    return BD
-def _npu_max_axis_chunks(grid_dim0: int, batch: int = 1) -> int:
-    denom = grid_dim0 * batch
-    if denom > _NPU_MAX_TRITON_GRID:
-        raise RuntimeError(
-            f'Ascend Triton grid dim0*batch={denom} exceeds {_NPU_MAX_TRITON_GRID}',
-        )
-    return max(1, _NPU_MAX_TRITON_GRID // max(denom, 1))
-def _npu_tile_config(
-    T: int,
-    BT: int,
-    D: int,
-    dtype: torch.dtype,
-    initial_state: torch.Tensor | None,
-) -> tuple[int, int, int]:
-    BT = _npu_chunk_size(T, BT)
-    BD = 16
-    if D >= 8192:
-        BD = 8
-        BT = min(BT, 8)
-    elif D >= 1024:
-        # BD=4 overflows Ascend UB on large-D forward; cap BT to limit NT.
-        BD = 8
-        BT = min(BT, 32)
-    elif D >= 512:
-        BD = 8
-    if dtype == torch.float16 and initial_state is not None:
-        BD = min(BD, 8)
-    if dtype == torch.bfloat16 and T <= 16:
-        BD = 8
-    return BD, BT, STATIC_WARPS
-def _npu_bwd_tile_config(
-    T: int,
-    BT: int,
-    D: int,
-    dtype: torch.dtype,
-    initial_state: torch.Tensor | None,
-) -> tuple[int, int, int]:
-    BT = _npu_chunk_size(T, BT)
-    BD = 16
-    if initial_state is not None:
-        BD = min(BD, 8)
-        BT = min(BT, 32)
-    if D >= 2048:
-        BD = 8
-        BT = min(BT, 8)
-    elif D >= 1024:
-        BD = 8
-        BT = min(BT, 16)
-    elif D >= 512:
-        BD = 8
-        BT = min(BT, 32)
-    if dtype == torch.bfloat16 and T <= 16:
-        BD = 8
-        BT = 32
-    return BD, BT, STATIC_WARPS
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['weight'] is not None,
-    'HAS_BIAS': lambda args: args['bias'] is not None,
-    'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None,
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
-@triton.jit
-def causal_conv1d_fwd_kernel(
-    x,
-    y,
-    weight,
-    bias,
-    cu_seqlens,
-    initial_state,
-    chunk_indices,
-    B,
-    T,
-    stride_x_n,
-    stride_x_t,
-    stride_x_d,
-    stride_y_n,
-    stride_y_t,
-    stride_y_d,
-    D: tl.constexpr,
-    W: tl.constexpr,
-    BT: tl.constexpr,
-    BW: tl.constexpr,
-    BD: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    USE_INITIAL_STATE: tl.constexpr,
-    IS_VARLEN: tl.constexpr,
-    CHUNK_OFFSET: tl.constexpr,
-):
-    i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
-        T = eos - bos
-        p_x = x + bos * stride_x_t
-        p_y = y + bos * stride_y_t
-    else:
-        i_n = i_b
-        i_t = i_t + CHUNK_OFFSET
-        bos = (i_b * T).to(tl.int64)
-        p_x = x + tl.cast(i_b, tl.int64) * stride_x_n
-        p_y = y + tl.cast(i_b, tl.int64) * stride_y_n
-    o_d = i_d * BD + tl.arange(0, BD)
-    o_w = tl.arange(0, BW) + W - BW
-    m_d = o_d < D
-    m_w = o_w >= 0
-    if HAS_WEIGHT:
-        b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0).to(tl.float32)
-    o_t = i_t * BT + tl.arange(0, BT)
-    m_t = (o_t >= 0) & (o_t < T)
-    b_y = tl.zeros((BT, BD), dtype=tl.float32)
-    for i_w in tl.static_range(-W + 1, 1):
-        o_x = o_t + i_w
-        m_x = ((o_x >= 0) & (o_x < T))[:, None] & m_d[None, :]
-        b_yi = tl.load(
-            p_x + o_x[:, None] * stride_x_t + o_d[None, :] * stride_x_d,
-            mask=m_x,
-            other=0,
-        ).to(tl.float32)
-        if USE_INITIAL_STATE:
-            m_c = ((o_x + W >= 0) & (o_x < 0))[:, None] & m_d[None, :]
-            b_yi += tl.load(
-                initial_state + i_n * D * W + o_d[None, :] * W + (o_x + W)[:, None],
-                mask=m_c,
-                other=0,
-            ).to(tl.float32)
-        if HAS_WEIGHT:
-            b_yi = b_yi * tl.sum(b_w * (o_w == (i_w + W - 1)), 1)[None, :]
-        b_y += b_yi
-    if HAS_BIAS:
-        b_y += tl.load(bias + o_d, mask=m_d).to(tl.float32)[None, :]
-    tl.store(
-        p_y + o_t[:, None] * stride_y_t + o_d[None, :] * stride_y_d,
-        tl.cast(b_y, dtype=y.dtype.element_ty, fp_downcast_rounding='rtne'),
-        mask=m_t[:, None] & m_d[None, :],
-    )
-@triton.jit
-def _silu_kernel(
-    x_ptr,
-    y_ptr,
-    n_elements,
-    ELEM_OFFSET: tl.constexpr,
-    BLOCK: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * BLOCK + tl.arange(0, BLOCK) + ELEM_OFFSET
-    mask = offs < n_elements
-    x = tl.load(x_ptr + offs, mask=mask, other=0.).to(tl.float32)
-    y = x * tl.sigmoid(x)
-    tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty), mask=mask)
-@triton.jit
-def _add_kernel(
-    a_ptr,
-    b_ptr,
-    out_ptr,
-    n_elements,
-    ELEM_OFFSET: tl.constexpr,
-    BLOCK: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * BLOCK + tl.arange(0, BLOCK) + ELEM_OFFSET
-    mask = offs < n_elements
-    a = tl.load(a_ptr + offs, mask=mask, other=0.).to(tl.float32)
-    b = tl.load(b_ptr + offs, mask=mask, other=0.).to(tl.float32)
-    tl.store(out_ptr + offs, (a + b).to(out_ptr.dtype.element_ty), mask=mask)
-def _launch_silu(y: torch.Tensor) -> torch.Tensor:
-    y = y.contiguous()
-    out = torch.zeros_like(y)
-    n = y.numel()
-    for grid, elem_off in _elementwise_launch_iters(n):
-        _silu_kernel[(grid,)](
-            y, out, n,
-            ELEM_OFFSET=elem_off,
-            BLOCK=_ELEM_BLOCK,
-            num_warps=STATIC_WARPS,
-        )
-    return out
-def _launch_add(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    a = a.contiguous()
-    b = b.contiguous()
-    out = torch.zeros_like(a)
-    n = a.numel()
-    for grid, elem_off in _elementwise_launch_iters(n):
-        _add_kernel[(grid,)](
-            a, b, out, n,
-            ELEM_OFFSET=elem_off,
-            BLOCK=_ELEM_BLOCK,
-            num_warps=STATIC_WARPS,
-        )
-    return out
-@triton.jit
-def _silu_bwd_kernel(
-    y_ptr,
-    dy_ptr,
-    out_ptr,
-    stride_y_n,
-    stride_y_t,
-    stride_y_d,
-    stride_dy_n,
-    stride_dy_t,
-    stride_dy_d,
-    stride_out_n,
-    stride_out_t,
-    stride_out_d,
-    B,
-    T,
-    D,
-    ELEM_OFFSET: tl.constexpr,
-    BLOCK: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * BLOCK + tl.arange(0, BLOCK) + ELEM_OFFSET
-    n_elements = B * T * D
-    mask = offs < n_elements
-    rem = offs % D
-    d = rem
-    rem = (offs - d) // D
-    t = rem % T
-    b = rem // T
-    y_off = b * stride_y_n + t * stride_y_t + d * stride_y_d
-    dy_off = b * stride_dy_n + t * stride_dy_t + d * stride_dy_d
-    out_off = b * stride_out_n + t * stride_out_t + d * stride_out_d
-    y = tl.load(y_ptr + y_off, mask=mask, other=0.).to(tl.float32)
-    dy = tl.load(dy_ptr + dy_off, mask=mask, other=0.).to(tl.float32)
-    s = tl.sigmoid(y)
-    out = dy * s * (1.0 + y * (1.0 - s))
-    tl.store(out_ptr + out_off, out.to(out_ptr.dtype.element_ty), mask=mask)
-def _launch_silu_bwd(y_pre: torch.Tensor, dy: torch.Tensor) -> torch.Tensor:
-    out = torch.zeros_like(dy, memory_format=torch.contiguous_format)
-    B, T, D = dy.shape
-    n = B * T * D
-    sy_n, sy_t, sy_d = y_pre.stride()
-    sdy_n, sdy_t, sdy_d = dy.stride()
-    so_n, so_t, so_d = out.stride()
-    for grid, elem_off in _elementwise_launch_iters(n):
-        _silu_bwd_kernel[(grid,)](
-            y_pre, dy, out,
-            sy_n, sy_t, sy_d,
-            sdy_n, sdy_t, sdy_d,
-            so_n, so_t, so_d,
-            B, T, D,
-            ELEM_OFFSET=elem_off,
-            BLOCK=_ELEM_BLOCK,
-            num_warps=STATIC_WARPS,
-        )
-    return out
-def _postprocess_fwd(
-    y: torch.Tensor,
-    residual: torch.Tensor | None,
-    activation: str | None,
-) -> torch.Tensor:
-    if activation in ('swish', 'silu'):
-        y = _launch_silu(y)
-    if residual is not None:
-        if residual.stride() != y.stride():
-            residual = residual.contiguous()
-        y = _launch_add(y, residual)
-    return y
-def _use_seq_bwd(
-    T: int,
-    dtype: torch.dtype,
-    initial_state: torch.Tensor | None,
-    dht: torch.Tensor | None,
-    cu_seqlens: torch.Tensor | None,
-) -> bool:
-    return (
-        cu_seqlens is None
-        and initial_state is None
-        and dht is None
-        and dtype == torch.bfloat16
-        and T <= 16
-    )
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['dw'] is not None,
-    'HAS_BIAS': lambda args: args['db'] is not None,
-})
-@triton.jit
-def causal_conv1d_bwd_seq_kernel(
-    x,
-    weight,
-    dy,
-    dx,
-    dw,
-    db,
-    stride_x_n,
-    stride_x_t,
-    stride_x_d,
-    stride_dx_n,
-    stride_dx_t,
-    stride_dx_d,
-    stride_dy_n,
-    stride_dy_t,
-    stride_dy_d,
-    B,
-    TC: tl.constexpr,
-    D: tl.constexpr,
-    W: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    BLOCK: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * BLOCK + tl.arange(0, BLOCK)
-    n_elements = B * TC * D
-    mask = offs < n_elements
-    d = offs % D
-    tmp = offs // D
-    t = tmp % TC
-    b = tmp // TC
-    b_dx = tl.zeros((BLOCK,), dtype=tl.float32)
-    for i_w in tl.static_range(0, W):
-        t_dy = t + i_w
-        dy_off = b * stride_dy_n + t_dy * stride_dy_t + d * stride_dy_d
-        b_dy = tl.load(dy + dy_off, mask=mask & (t_dy < TC), other=0.).to(tl.float32)
-        if HAS_WEIGHT:
-            w_idx = W - i_w - 1
-            b_w = tl.load(weight + d * W + w_idx, mask=mask, other=0.).to(tl.float32)
-            b_dx += b_dy * b_w
-        else:
-            b_dx += b_dy
-    dx_off = b * stride_dx_n + t * stride_dx_t + d * stride_dx_d
-    tl.store(dx + dx_off, b_dx.to(dx.dtype.element_ty), mask=mask)
-    if HAS_WEIGHT:
-        x_off = b * stride_x_n + t * stride_x_t + d * stride_x_d
-        b_x = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
-        i_tg = b * TC + t
-        for i_w in tl.static_range(0, W):
-            t_dy = t + i_w
-            dy_off = b * stride_dy_n + t_dy * stride_dy_t + d * stride_dy_d
-            b_dy = tl.load(dy + dy_off, mask=mask & (t_dy < TC), other=0.).to(tl.float32)
-            w_idx = W - i_w - 1
-            tl.store(
-                dw + (i_tg * D + d) * W + w_idx,
-                (b_dy * b_x).to(dw.dtype.element_ty),
-                mask=mask,
-            )
-    if HAS_BIAS:
-        i_tg = b * TC + t
-        dy_off = b * stride_dy_n + t * stride_dy_t + d * stride_dy_d
-        b_dy0 = tl.load(dy + dy_off, mask=mask, other=0.)
-        tl.store(db + i_tg * D + d, b_dy0.to(db.dtype.element_ty), mask=mask)
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['dw'] is not None,
-    'HAS_BIAS': lambda args: args['db'] is not None,
-    'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None,
-    'USE_FINAL_STATE': lambda args: args['dht'] is not None,
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
-@triton.jit
-def causal_conv1d_bwd_kernel(
-    x,
-    weight,
-    initial_state,
-    dht,
-    dy,
-    dx,
-    dw,
-    db,
-    cu_seqlens,
-    chunk_indices,
-    B,
-    T,
-    stride_x_n,
-    stride_x_t,
-    stride_x_d,
-    stride_dx_n,
-    stride_dx_t,
-    stride_dx_d,
-    stride_dy_n,
-    stride_dy_t,
-    stride_dy_d,
-    D: tl.constexpr,
-    W: tl.constexpr,
-    BT: tl.constexpr,
-    BW: tl.constexpr,
-    BD: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    USE_INITIAL_STATE: tl.constexpr,
-    USE_FINAL_STATE: tl.constexpr,
-    IS_VARLEN: tl.constexpr,
-    CHUNK_OFFSET: tl.constexpr,
-    NT: tl.constexpr,
-):
-    i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    if IS_VARLEN:
-        i_tg = i_t
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
-        T = eos - bos
-        p_x = x + bos * stride_x_t
-        p_dy = dy + bos * stride_dy_t
-        p_dx = dx + bos * stride_dx_t
-    else:
-        i_t = i_t + CHUNK_OFFSET
-        i_tg = i_b * NT + i_t
-        i_n = i_b
-        p_x = x + tl.cast(i_b, tl.int64) * stride_x_n
-        p_dy = dy + tl.cast(i_b, tl.int64) * stride_dy_n
-        p_dx = dx + tl.cast(i_b, tl.int64) * stride_dx_n
-    o_d = i_d * BD + tl.arange(0, BD)
-    o_w = tl.arange(0, BW) + W - BW
-    m_d = o_d < D
-    m_w = o_w >= 0
-    o_t = i_t * BT + tl.arange(0, BT)
-    m_t = (o_t >= 0) & (o_t < T)
-    b_x = tl.zeros((BT, BD), dtype=tl.float32)
-    if HAS_WEIGHT:
-        b_x = tl.load(
-            p_x + o_t[:, None] * stride_x_t + o_d[None, :] * stride_x_d,
-            mask=m_t[:, None] & m_d[None, :],
-            other=0,
-        ).to(tl.float32)
-        b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0).to(tl.float32)
-    b_dx = tl.zeros((BT, BD), dtype=tl.float32)
-    if HAS_BIAS:
-        b_db = tl.zeros((BD,), dtype=tl.float32)
-    for i_w in tl.static_range(0, W):
-        o_dy = o_t + i_w
-        m_dy = ((o_dy >= 0) & (o_dy < T))[:, None] & m_d[None, :]
-        b_dy = tl.load(
-            p_dy + o_dy[:, None] * stride_dy_t + o_d[None, :] * stride_dy_d,
-            mask=m_dy,
-            other=0,
-        ).to(tl.float32)
-        if HAS_WEIGHT:
-            b_wdy = b_dy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1)[None, :]
-            b_dw = tl.sum(b_dy * b_x, 0)
-            if USE_INITIAL_STATE:
-                mask_head_rows = (o_t < i_w) & (o_t < T)
-                b_dy_head = tl.load(
-                    p_dy + o_t[:, None] * stride_dy_t + o_d[None, :] * stride_dy_d,
-                    mask=(mask_head_rows[:, None] & m_d[None, :]),
-                    other=0.0,
-                ).to(tl.float32)
-                o_c = W - i_w + o_t
-                mask_c = (mask_head_rows & (o_c >= 1) & (o_c < W))
-                b_xc = tl.load(
-                    initial_state + i_n * D * W + o_d[None, :] * W + o_c[:, None],
-                    mask=(mask_c[:, None] & m_d[None, :]),
-                    other=0.0,
-                ).to(tl.float32)
-                b_dw += tl.sum(b_dy_head * b_xc, 0)
-            tl.store(dw + i_tg * D * W + o_d * W + W - i_w - 1, b_dw.to(dw.dtype.element_ty), mask=m_d)
-        else:
-            b_wdy = b_dy
-        if HAS_BIAS and i_w == 0:
-            b_db += tl.sum(b_dy, 0)
-        b_dx += b_wdy
-    if HAS_BIAS:
-        b_db = tl.cast(b_db, dtype=db.dtype.element_ty, fp_downcast_rounding='rtne')
-        tl.store(db + i_tg * D + o_d, b_db, mask=m_d)
-    if USE_FINAL_STATE:
-        if i_t * BT + BT >= T - W:
-            start_tok = T - (W - 1)
-            offset = i_t * BT + tl.arange(0, BT)
-            tok_idx = offset - start_tok
-            mask = (offset >= start_tok) & (offset < T)
-            w_idx = 1 + tok_idx
-            dht_off = i_n * D * W + o_d[None, :] * W + w_idx[:, None]
-            b_dht = tl.load(dht + dht_off, mask=mask[:, None] & m_d[None, :], other=0.).to(tl.float32)
-            b_dx += b_dht
-    tl.store(
-        p_dx + o_t[:, None] * stride_dx_t + o_d[None, :] * stride_dx_d,
-        tl.cast(b_dx, dtype=dx.dtype.element_ty, fp_downcast_rounding='rtne'),
-        mask=m_t[:, None] & m_d[None, :],
-    )
-@triton.heuristics({
-    'USE_ACTIVATION': lambda args: args['y'] is not None,
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
-@triton.jit
-def compute_dh0_kernel(
-    dy,
-    y,
-    weight,
-    dh0,
-    cu_seqlens,
-    stride_dy_n,
-    stride_dy_t,
-    stride_dy_d,
-    stride_y_n,
-    stride_y_t,
-    stride_y_d,
-    T,
-    D: tl.constexpr,
-    W: tl.constexpr,
-    BD: tl.constexpr,
-    USE_ACTIVATION: tl.constexpr,
-    IS_VARLEN: tl.constexpr,
-    CHUNK_OFFSET: tl.constexpr,
-):
-    i_d, i_n = tl.program_id(0), tl.program_id(1) + CHUNK_OFFSET
-    if IS_VARLEN:
-        bos = tl.load(cu_seqlens + i_n).to(tl.int64)
-        eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
-        seq_len = eos - bos
-        dy_base = dy + bos * stride_dy_t
-    else:
-        seq_len = T
-        dy_base = dy + tl.cast(i_n, tl.int64) * stride_dy_n
-    o_d = i_d * BD + tl.arange(0, BD)
-    m_d = o_d < D
-    for i_w in tl.static_range(1, W):
-        b_dh0 = tl.zeros([BD], dtype=tl.float32)
-        for t in tl.static_range(0, W - 1):
-            if t < i_w:
-                w_idx = i_w - 1 - t
-                p_dy = dy_base + t * stride_dy_t + o_d * stride_dy_d
-                m_t = (t < seq_len) & m_d
-                b_dy = tl.load(p_dy, mask=m_t, other=0).to(tl.float32)
-                if USE_ACTIVATION:
-                    if IS_VARLEN:
-                        p_y = y + bos * stride_y_t + t * stride_y_t + o_d * stride_y_d
-                    else:
-                        p_y = y + tl.cast(i_n, tl.int64) * stride_y_n + t * stride_y_t + o_d * stride_y_d
-                    b_y = tl.load(p_y, mask=m_t, other=0).to(tl.float32)
-                    b_ys = tl.sigmoid(b_y)
-                    b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
-                b_w_col = tl.load(weight + o_d * W + w_idx, mask=m_d, other=0).to(tl.float32)
-                b_dh0 += tl.where(m_t, b_dy * b_w_col, 0)
-        p_dh0 = dh0 + i_n * D * W + o_d * W + i_w
-        tl.store(p_dh0, b_dh0.to(dh0.dtype.element_ty), mask=m_d)
-@triton.heuristics({
-    'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None,
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
-@triton.jit
-def causal_conv1d_states_fwd_kernel(
-    x,
-    initial_state,
-    final_state,
-    cu_seqlens,
-    T,
-    D,
-    W,
-    stride_x_n,
-    stride_x_t,
-    stride_x_d,
-    BD: tl.constexpr,
-    BW: tl.constexpr,
-    USE_INITIAL_STATE: tl.constexpr,
-    IS_VARLEN: tl.constexpr,
-    CHUNK_OFFSET: tl.constexpr,
-):
-    i_d, i_n = tl.program_id(0), tl.program_id(1) + CHUNK_OFFSET
-    o_d = i_d * BD + tl.arange(0, BD)
-    m_d = o_d < D
-    if IS_VARLEN:
-        bos = tl.load(cu_seqlens + i_n).to(tl.int64)
-        eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
-        seq_len = (eos - bos).to(tl.int32)
-        p_x = x + bos * stride_x_t
-    else:
-        seq_len = T
-        p_x = x + tl.cast(i_n, tl.int64) * stride_x_n
-    o_w = W - BW + tl.arange(0, BW)
-    m_w = o_w >= 0
-    o_t = seq_len - BW + tl.arange(0, BW)
-    m_t = (o_t >= 0) & (o_t < seq_len)
-    b_x = tl.load(
-        p_x + o_t[:, None] * stride_x_t + o_d[None, :] * stride_x_d,
-        mask=m_t[:, None] & m_d[None, :],
-        other=0,
-    ).to(tl.float32)
-    if USE_INITIAL_STATE:
-        if seq_len < BW:
-            o_c = W - (BW - seq_len) + tl.arange(0, BW)
-            m_c = (o_c >= 0) & (o_c < W)
-            b_cache = tl.load(
-                initial_state + i_n * D * W + o_d[None, :] * W + o_c[:, None],
-                mask=m_d[None, :] & m_c[:, None],
-                other=0,
-            ).to(tl.float32)
-            b_x += b_cache
-    p_final = final_state + tl.cast(i_n, tl.int64) * D * W + o_d[:, None] * W + o_w[None, :]
-    tl.store(p_final, tl.trans(b_x).to(final_state.dtype.element_ty), mask=m_d[:, None] & m_w[None, :])
-@triton.heuristics({
-    'HAS_WEIGHT': lambda args: args['weight'] is not None,
-    'HAS_BIAS': lambda args: args['bias'] is not None,
-})
-@triton.jit
-def causal_conv1d_update_kernel(
-    x,
-    cache,
-    y,
-    weight,
-    bias,
-    stride_x_n,
-    stride_x_d,
-    stride_y_n,
-    stride_y_d,
-    D: tl.constexpr,
-    W: tl.constexpr,
-    BD: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    CHUNK_OFFSET: tl.constexpr,
-):
-    i_d, i_n = tl.program_id(0), tl.program_id(1) + CHUNK_OFFSET
-    o_d = i_d * BD + tl.arange(0, BD)
-    m_d = o_d < D
-    b_x = tl.load(x + i_n * stride_x_n + o_d * stride_x_d, mask=m_d, other=0).to(tl.float32)
-    b_y = tl.zeros((BD,), dtype=tl.float32)
-    for iw in tl.static_range(0, W):
-        if iw < W - 1:
-            b_c = tl.load(cache + i_n * D * W + o_d * W + (iw + 1), mask=m_d, other=0).to(tl.float32)
-        else:
-            b_c = b_x
-        tl.store(
-            cache + i_n * D * W + o_d * W + iw,
-            tl.cast(b_c, dtype=cache.dtype.element_ty, fp_downcast_rounding='rtne'),
-            mask=m_d,
-        )
-        if HAS_WEIGHT:
-            b_y += b_c * tl.load(weight + o_d * W + iw, mask=m_d, other=0).to(tl.float32)
-        else:
-            b_y += b_c
-    if HAS_BIAS:
-        b_y += tl.load(bias + o_d, mask=m_d)
-    tl.store(
-        y + i_n * stride_y_n + o_d * stride_y_d,
-        tl.cast(b_y, dtype=y.dtype.element_ty, fp_downcast_rounding='rtne'),
-        mask=m_d,
-    )
-def _postprocess_update(
-    y: torch.Tensor,
-    residual: torch.Tensor | None,
-    activation: str | None,
-) -> torch.Tensor:
-    if activation in ('swish', 'silu'):
-        y = _launch_silu(y)
-    if residual is not None:
-        if residual.stride() != y.stride():
-            residual = residual.contiguous()
-        y = _launch_add(y, residual)
-    return y
-def _launch_fwd_core(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    initial_state: torch.Tensor | None,
-    cu_seqlens: torch.LongTensor | None,
-    chunk_indices: torch.LongTensor | None,
-    B: int,
-    T: int,
-    D: int,
-    W: int,
-    BT: int,
-    BD: int | None = None,
-    num_warps: int | None = None,
-) -> torch.Tensor:
-    if BD is None or num_warps is None:
-        BD, BT, num_warps = _npu_tile_config(T, BT, D, x.dtype, initial_state)
-    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
-    BD = _clamp_bd_for_grid(B, NT, D, BD)
-    BW = triton.next_power_of_2(W)
-    stride_x_n, stride_x_t, stride_x_d = x.stride()
-    y = torch.zeros_like(x, memory_format=torch.contiguous_format)
-    stride_y_n, stride_y_t, stride_y_d = y.stride()
-    max_nt = _npu_max_axis_chunks(triton.cdiv(D, BD), B)
-    kernel_kwargs = dict(
-        x=x,
-        y=y,
-        weight=weight,
-        bias=bias,
-        cu_seqlens=cu_seqlens,
-        initial_state=initial_state,
-        B=B,
-        T=T,
-        D=D,
-        W=W,
-        BT=BT,
-        BW=BW,
-        BD=BD,
-        stride_x_n=stride_x_n,
-        stride_x_t=stride_x_t,
-        stride_x_d=stride_x_d,
-        stride_y_n=stride_y_n,
-        stride_y_t=stride_y_t,
-        stride_y_d=stride_y_d,
-        num_warps=num_warps,
-    )
-    for nt_off in range(0, NT, max_nt):
-        nt_len = min(max_nt, NT - nt_off)
-        grid = (triton.cdiv(D, BD), nt_len, B)
-        if cu_seqlens is not None:
-            kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len]
-            kernel_kwargs['CHUNK_OFFSET'] = 0
-        else:
-            kernel_kwargs['chunk_indices'] = chunk_indices
-            kernel_kwargs['CHUNK_OFFSET'] = nt_off
-        causal_conv1d_fwd_kernel[grid](**kernel_kwargs)
-    return y
-@input_guard(no_guard_contiguous=['x'])
-def causal_conv1d_fwd_npu(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    residual: torch.Tensor,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool = False,
-    activation: str | None = None,
-    cu_seqlens: torch.LongTensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    chunk_indices: torch.LongTensor | None = None,
-    BT: int = 64,
-    layout_fallback: bool = False,
-):
-    del layout_fallback
-    shape = x.shape
-    if x.shape[-1] != weight.shape[0]:
-        x = rearrange(x, 'b t ... -> b t (...)')
-    B, T, D = x.shape[0], x.shape[1], weight.shape[0]
-    W = weight.shape[1]
-    BD, BT, num_warps = _npu_tile_config(T, BT, D, x.dtype, initial_state)
-    if cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
-    y = _launch_fwd_core(
-        x, weight, bias, initial_state, cu_seqlens, chunk_indices, B, T, D, W, BT, BD, num_warps,
-    )
-    y = _postprocess_fwd(y, residual, activation)
-    final_state = None
-    if output_final_state:
-        final_state = causal_conv1d_update_states_npu(
-            x=x,
-            state_len=W,
-            initial_state=initial_state,
-            cu_seqlens=cu_seqlens,
-        )
-    return y.view(shape), final_state
-def causal_conv1d_bwd_npu(
-    x: torch.Tensor,
-    dy: torch.Tensor,
-    dht: torch.Tensor,
-    weight: torch.Tensor | None = None,
-    bias: torch.Tensor | None = None,
-    residual: torch.Tensor | None = None,
-    initial_state: torch.Tensor | None = None,
-    activation: str | None = None,
-    cu_seqlens: torch.Tensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    chunk_indices: torch.LongTensor | None = None,
-    BT: int = 64,
-    layout_fallback: bool = False,
-):
-    del layout_fallback
-    shape = x.shape
-    if x.shape[-1] != weight.shape[0]:
-        x = rearrange(x, 'b t ... -> b t (...)')
-    B, T, D = x.shape
-    W = weight.shape[1] if weight is not None else None
-    BD, BT, num_warps = _npu_bwd_tile_config(T, BT, D, x.dtype, initial_state)
-    if cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
-    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
-    BD = _clamp_bd_for_grid(B, NT, D, BD)
-    BW = triton.next_power_of_2(W)
-    dr = dy if residual is not None else None
-    dy_conv = dy
-    y_pre = None
-    if activation in ('swish', 'silu'):
-        BD_f, BT_f, nw_f = _npu_tile_config(T, BT, D, x.dtype, initial_state)
-        chunk_indices_f = chunk_indices
-        if cu_seqlens is not None:
-            chunk_indices_f = prepare_chunk_indices(cu_seqlens, BT_f, cu_seqlens_cpu=cu_seqlens_cpu)
-        y_pre = _launch_fwd_core(
-            x, weight, bias, initial_state, cu_seqlens, chunk_indices_f,
-            B, T, D, W, BT_f, BD_f, nw_f,
-        )
-        dy_conv = _launch_silu_bwd(y_pre, dy)
-    stride_x_n, stride_x_t, stride_x_d = x.stride()
-    use_seq = _use_seq_bwd(T, x.dtype, initial_state, dht, cu_seqlens)
-    stride_dy_n, stride_dy_t, stride_dy_d = dy_conv.stride()
-    dx = torch.zeros_like(x)
-    stride_dx_n, stride_dx_t, stride_dx_d = dx.stride()
-    if use_seq:
-        block = 1024
-        dw = weight.new_empty(B * T, *weight.shape, dtype=torch.float) if weight is not None else None
-        db = bias.new_empty(B * T, *bias.shape, dtype=torch.float) if bias is not None else None
-        grid = (triton.cdiv(B * T * D, block),)
-        causal_conv1d_bwd_seq_kernel[grid](
-            x=x,
-            weight=weight,
-            dy=dy_conv,
-            dx=dx,
-            dw=dw,
-            db=db,
-            stride_x_n=stride_x_n,
-            stride_x_t=stride_x_t,
-            stride_x_d=stride_x_d,
-            stride_dx_n=stride_dx_n,
-            stride_dx_t=stride_dx_t,
-            stride_dx_d=stride_dx_d,
-            stride_dy_n=stride_dy_n,
-            stride_dy_t=stride_dy_t,
-            stride_dy_d=stride_dy_d,
-            B=B,
-            TC=T,
-            D=D,
-            W=W,
-            BLOCK=block,
-            num_warps=STATIC_WARPS,
-        )
-    else:
-        if not dy_conv.is_contiguous():
-            dy_conv = dy_conv.contiguous()
-        stride_dy_n, stride_dy_t, stride_dy_d = dy_conv.stride()
-        dw = weight.new_empty(B * NT, *weight.shape, dtype=torch.float) if weight is not None else None
-        db = bias.new_empty(B * NT, *bias.shape, dtype=torch.float) if bias is not None else None
-        max_nt = _npu_max_axis_chunks(triton.cdiv(D, BD), B)
-        kernel_kwargs = dict(
-            x=x,
-            weight=weight,
-            initial_state=initial_state,
-            dht=dht,
-            dy=dy_conv,
-            dx=dx,
-            cu_seqlens=cu_seqlens,
-            B=B,
-            T=T,
-            D=D,
-            W=W,
-            BT=BT,
-            BW=BW,
-            BD=BD,
-            stride_x_n=stride_x_n,
-            stride_x_t=stride_x_t,
-            stride_x_d=stride_x_d,
-            stride_dx_n=stride_dx_n,
-            stride_dx_t=stride_dx_t,
-            stride_dx_d=stride_dx_d,
-            stride_dy_n=stride_dy_n,
-            stride_dy_t=stride_dy_t,
-            stride_dy_d=stride_dy_d,
-            num_warps=num_warps,
-            NT=NT,
-        )
-        for nt_off in range(0, NT, max_nt):
-            nt_len = min(max_nt, NT - nt_off)
-            grid = (triton.cdiv(D, BD), nt_len, B)
-            if cu_seqlens is not None:
-                kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len]
-                kernel_kwargs['CHUNK_OFFSET'] = 0
-                kernel_kwargs['dw'] = dw[nt_off:nt_off + nt_len] if weight is not None else None
-                kernel_kwargs['db'] = db[nt_off:nt_off + nt_len] if bias is not None else None
-            else:
-                kernel_kwargs['chunk_indices'] = chunk_indices
-                kernel_kwargs['CHUNK_OFFSET'] = nt_off
-                kernel_kwargs['dw'] = dw
-                kernel_kwargs['db'] = db
-            causal_conv1d_bwd_kernel[grid](**kernel_kwargs)
-    if weight is not None:
-        dw = dw.sum(0).to(weight)
-    if bias is not None:
-        db = db.sum(0).to(bias)
-    dh0 = None
-    if initial_state is not None:
-        dh0 = compute_dh0_npu(
-            dy=dy,
-            y=y_pre,
-            weight=weight,
-            initial_state=initial_state,
-            activation=activation,
-            cu_seqlens=cu_seqlens,
-        )
-    return dx.view(shape), dw, db, dr, dh0
-def compute_dh0_npu(
-    dy: torch.Tensor,
-    y: torch.Tensor | None,
-    weight: torch.Tensor,
-    initial_state: torch.Tensor,
-    activation: str | None,
-    cu_seqlens: torch.Tensor | None,
-) -> torch.Tensor:
-    D, W = weight.shape
-    N = initial_state.shape[0]
-    T = dy.shape[1]
-    BD = 8 if dy.dtype == torch.float16 and activation in ('swish', 'silu') else 16
-    dh0 = torch.zeros_like(initial_state)
-    stride_dy_n = dy.stride(0)
-    stride_dy_t = dy.stride(1)
-    stride_dy_d = dy.stride(2) if dy.dim() == 3 else dy.stride(-1)
-    stride_y_n = stride_y_t = stride_y_d = 0
-    if y is not None:
-        stride_y_n = y.stride(0)
-        stride_y_t = y.stride(1)
-        stride_y_d = y.stride(2) if y.dim() == 3 else y.stride(-1)
-    max_n = _npu_max_axis_chunks(triton.cdiv(D, BD))
-    kernel_kwargs = dict(
-        dy=dy,
-        y=y if activation in ('swish', 'silu') else None,
-        weight=weight,
-        dh0=dh0,
-        cu_seqlens=cu_seqlens,
-        stride_dy_n=stride_dy_n,
-        stride_dy_t=stride_dy_t,
-        stride_dy_d=stride_dy_d,
-        stride_y_n=stride_y_n,
-        stride_y_t=stride_y_t,
-        stride_y_d=stride_y_d,
-        T=T,
-        D=D,
-        W=W,
-        BD=BD,
-        num_warps=STATIC_WARPS,
-    )
-    for n_off in range(0, N, max_n):
-        n_len = min(max_n, N - n_off)
-        kernel_kwargs['CHUNK_OFFSET'] = n_off
-        compute_dh0_kernel[(triton.cdiv(D, BD), n_len)](**kernel_kwargs)
-    return dh0
-@input_guard(no_guard_contiguous=['x'])
-def causal_conv1d_update_states_npu(
-    x: torch.Tensor,
-    state_len: int,
-    initial_state: torch.Tensor | None = None,
-    cu_seqlens: torch.Tensor | None = None,
-    layout_fallback: bool = False,
-) -> torch.Tensor:
-    del layout_fallback
-    if cu_seqlens is not None:
-        N = len(cu_seqlens) - 1
-        if x.dim() == 2:
-            stride_x_n = 0
-            stride_x_t, stride_x_d = x.stride()
-            T = x.shape[0]
-        else:
-            stride_x_n = x.stride(0)
-            stride_x_t, stride_x_d = x.stride(1), x.stride(2)
-            T = x.shape[1]
-        D = x.shape[-1]
-    else:
-        B, T, D = x.shape
-        N = B
-        stride_x_n, stride_x_t, stride_x_d = x.stride()
-    W = state_len
-    final_state = torch.empty(N, D, W, dtype=x.dtype, device=x.device)
-    BD = min(triton.next_power_of_2(D), 16)
-    BW = triton.next_power_of_2(W)
-    grid_dim0 = triton.cdiv(D, BD)
-    max_n = _npu_max_axis_chunks(grid_dim0)
-    kernel_kwargs = dict(
-        x=x,
-        initial_state=initial_state,
-        final_state=final_state,
-        cu_seqlens=cu_seqlens,
-        T=T,
-        D=D,
-        W=W,
-        stride_x_n=stride_x_n,
-        stride_x_t=stride_x_t,
-        stride_x_d=stride_x_d,
-        BW=BW,
-        BD=BD,
-        num_warps=STATIC_WARPS,
-    )
-    for n_off in range(0, N, max_n):
-        n_len = min(max_n, N - n_off)
-        kernel_kwargs['CHUNK_OFFSET'] = n_off
-        causal_conv1d_states_fwd_kernel[(grid_dim0, n_len)](**kernel_kwargs)
-    return final_state
-@input_guard(no_guard_contiguous=['x'])
-def causal_conv1d_update_npu(
-    x: torch.Tensor,
-    cache: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    weight: torch.Tensor | None = None,
-    bias: torch.Tensor | None = None,
-    activation: str | None = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    shape = x.shape
-    if weight is not None and x.shape[-1] != weight.shape[0]:
-        x = rearrange(x, 'b t ... -> b t (...)')
-    D = x.shape[-1]
-    N = x.numel() // D
-    W = weight.shape[1] if weight is not None else None
-    BD = min(triton.next_power_of_2(D), 16)
-    if x.dim() == 2:
-        stride_x_n = x.stride(0)
-        stride_x_d = x.stride(1)
-    elif x.dim() == 3 and x.shape[0] == 1:
-        stride_x_n = x.stride(1)
-        stride_x_d = x.stride(2)
-    elif x.dim() == 3:
-        stride_x_n = x.stride(0)
-        stride_x_d = x.stride(2)
-    else:
-        raise ValueError(f"Unsupported input shape: {x.shape}")
-    y = torch.zeros_like(x, memory_format=torch.contiguous_format)
-    if y.dim() == 2:
-        stride_y_n, stride_y_d = y.stride(0), y.stride(1)
-    elif y.dim() == 3 and y.shape[0] == 1:
-        stride_y_n, stride_y_d = y.stride(1), y.stride(2)
-    elif y.dim() == 3:
-        stride_y_n, stride_y_d = y.stride(0), y.stride(2)
-    grid_dim0 = triton.cdiv(D, BD)
-    max_n = _npu_max_axis_chunks(grid_dim0)
-    kernel_kwargs = dict(
-        x=x,
-        cache=cache,
-        y=y,
-        weight=weight,
-        bias=bias,
-        stride_x_n=stride_x_n,
-        stride_x_d=stride_x_d,
-        stride_y_n=stride_y_n,
-        stride_y_d=stride_y_d,
-        D=D,
-        W=W,
-        BD=BD,
-        num_warps=STATIC_WARPS,
-    )
-    for n_off in range(0, N, max_n):
-        n_len = min(max_n, N - n_off)
-        kernel_kwargs['CHUNK_OFFSET'] = n_off
-        causal_conv1d_update_kernel[(grid_dim0, n_len)](**kernel_kwargs)
-    y = _postprocess_update(y, residual, activation)
-    return y.view(shape), cache

build/torch-cuda/modules/backends/triton_ascend/fused_cross_entropy.py DELETED Viewed

@@ -1,469 +0,0 @@
-# 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
-"""Fused cross-entropy kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import triton
-import triton.language as tl
-from triton.language.math import tanh
-from ....ops.utils.op import exp, log
-from ....utils import input_guard
-from ....utils.ascend_ub_manager import (
-    ASCEND_MAX_GRID_DIM,
-    compute_vocab_block_size,
-    iter_axis_launch_chunks,
-)
-# Cross-entropy fwd/bwd peak fp32 buffers along vocab dimension.
-_CE_FWD_MEM_MULT = 8.0
-_CE_BWD_MEM_MULT = 12.0
-@triton.heuristics({
-    "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0,
-})
-@triton.jit
-def cross_entropy_fwd_kernel(
-    loss_ptr,
-    lse_ptr,
-    z_loss_ptr,
-    logits_ptr,
-    labels_ptr,
-    label_smoothing,
-    logit_scale,
-    lse_square_scale,
-    logit_softcapping: tl.constexpr,
-    ignore_index,
-    total_classes,
-    class_start_idx,
-    n_cols,
-    n_rows,
-    logits_row_stride,
-    ROW_OFFSET,
-    BLOCK_SIZE: tl.constexpr,
-    HAS_SMOOTHING: tl.constexpr,
-    HAS_SOFTCAPPING: tl.constexpr,
-    SPLIT: tl.constexpr,
-):
-    row_idx = tl.program_id(0)
-    abs_row_idx = row_idx + ROW_OFFSET
-    col_block_idx = tl.program_id(1)
-    logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64)
-    col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    label_idx = tl.load(labels_ptr + row_idx)
-    logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf"))
-    logits = logits.to(tl.float32) * logit_scale
-    if HAS_SOFTCAPPING:
-        logits = logit_softcapping * tanh(logits / logit_softcapping)
-    max_logits = tl.max(logits, 0)
-    if HAS_SMOOTHING:
-        sum_logits = tl.sum(tl.where(col_offsets < n_cols, logits, 0.0), 0)
-    lse = log(tl.sum(exp(logits - max_logits), 0)) + max_logits
-    tl.store(lse_ptr + col_block_idx * n_rows + abs_row_idx, lse)
-    if label_idx == ignore_index:
-        loss = 0.0
-        z_loss = 0.0
-    else:
-        label_idx -= class_start_idx
-        if label_idx >= col_block_idx * BLOCK_SIZE and label_idx < min(
-            n_cols, (col_block_idx + 1) * BLOCK_SIZE,
-        ):
-            logits_label = tl.load(logits_ptr + label_idx).to(tl.float32) * logit_scale
-            if HAS_SOFTCAPPING:
-                logits_label = logit_softcapping * tanh(logits_label / logit_softcapping)
-            if HAS_SMOOTHING:
-                loss = (
-                    (lse if not SPLIT else 0.0)
-                    - label_smoothing * sum_logits / total_classes
-                    - (1 - label_smoothing) * logits_label
-                )
-            else:
-                loss = (lse if not SPLIT else 0.0) - logits_label
-        else:
-            if HAS_SMOOTHING:
-                loss = label_smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes)
-            else:
-                loss = 0.0
-        if not SPLIT:
-            z_loss = lse_square_scale * lse * lse
-            loss += z_loss
-        else:
-            z_loss = 0.0
-    tl.store(loss_ptr + col_block_idx * n_rows + abs_row_idx, loss)
-    if not SPLIT:
-        tl.store(z_loss_ptr + col_block_idx * n_rows + abs_row_idx, z_loss)
-@triton.heuristics({
-    "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0,
-})
-@triton.jit
-def cross_entropy_bwd_kernel(
-    dlogits_ptr,
-    dloss_ptr,
-    logits_ptr,
-    lse_ptr,
-    labels_ptr,
-    label_smoothing,
-    logit_scale,
-    lse_square_scale,
-    logit_softcapping: tl.constexpr,
-    ignore_index,
-    total_classes,
-    class_start_idx,
-    n_cols,
-    logits_row_stride,
-    dlogits_row_stride,
-    dloss_row_stride,
-    ROW_OFFSET,
-    BLOCK_SIZE: tl.constexpr,
-    HAS_SMOOTHING: tl.constexpr,
-    HAS_SOFTCAPPING: tl.constexpr,
-):
-    row_idx = tl.program_id(0)
-    abs_row_idx = row_idx + ROW_OFFSET
-    col_block_idx = tl.program_id(1)
-    logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64)
-    dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64)
-    col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    label_idx = tl.load(labels_ptr + row_idx)
-    if label_idx != ignore_index:
-        dloss = tl.load(dloss_ptr + abs_row_idx * dloss_row_stride)
-    else:
-        dloss = 0.0
-    logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to(
-        tl.float32,
-    ) * logit_scale
-    if HAS_SOFTCAPPING:
-        t = tanh(logits / logit_softcapping)
-        logits = logit_softcapping * t
-    lse = tl.load(lse_ptr + abs_row_idx)
-    probs = exp(logits - lse)
-    probs += 2.0 * lse_square_scale * lse * probs
-    label_idx -= class_start_idx
-    if HAS_SMOOTHING:
-        smooth_negative = label_smoothing / total_classes
-        probs = tl.where(col_offsets == label_idx, probs - (1 - label_smoothing), probs) - smooth_negative
-    else:
-        probs = tl.where(col_offsets == label_idx, probs - 1.0, probs)
-    if HAS_SOFTCAPPING:
-        probs = probs * (1.0 - t * t)
-    tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols)
-def _npu_block_size(n_cols: int, n_rows: int, is_backward: bool = False) -> tuple[int, int]:
-    memory_multiplier = _CE_BWD_MEM_MULT if is_backward else _CE_FWD_MEM_MULT
-    block_size = compute_vocab_block_size(n_cols, n_rows, memory_multiplier)
-    num_warps = 2 if block_size <= 2048 else 4
-    return block_size, num_warps
-def _launch_cross_entropy_fwd(
-    losses,
-    lse,
-    z_losses,
-    logits,
-    target,
-    *,
-    label_smoothing,
-    logit_scale,
-    lse_square_scale,
-    softcap_val,
-    ignore_index,
-    total_classes,
-    class_start_idx,
-    n_cols,
-    n_rows,
-    logits_stride,
-    BLOCK_SIZE,
-    has_softcapping,
-    num_warps,
-    split,
-):
-    n_splits = triton.cdiv(n_cols, BLOCK_SIZE)
-    for row_off, row_len in iter_axis_launch_chunks(n_rows, n_splits, max_grid=ASCEND_MAX_GRID_DIM):
-        cross_entropy_fwd_kernel[(row_len, n_splits)](
-            losses,
-            lse,
-            z_losses,
-            logits[row_off:row_off + row_len],
-            target[row_off:row_off + row_len],
-            label_smoothing,
-            logit_scale,
-            lse_square_scale,
-            softcap_val,
-            ignore_index,
-            total_classes,
-            class_start_idx,
-            n_cols,
-            n_rows,
-            logits_stride,
-            row_off,
-            BLOCK_SIZE=BLOCK_SIZE,
-            HAS_SOFTCAPPING=has_softcapping,
-            num_warps=num_warps,
-            SPLIT=split,
-        )
-def _launch_cross_entropy_bwd(
-    dlogits,
-    grad_losses,
-    logits,
-    lse,
-    target,
-    *,
-    label_smoothing,
-    logit_scale,
-    lse_square_scale,
-    softcap_val,
-    ignore_index,
-    total_classes,
-    class_start_idx,
-    n_cols,
-    n_rows,
-    logits_stride,
-    dlogits_stride,
-    grad_losses_stride,
-    BLOCK_SIZE,
-    has_softcapping,
-    num_warps,
-):
-    n_splits = triton.cdiv(n_cols, BLOCK_SIZE)
-    for row_off, row_len in iter_axis_launch_chunks(n_rows, n_splits, max_grid=ASCEND_MAX_GRID_DIM):
-        cross_entropy_bwd_kernel[(row_len, n_splits)](
-            dlogits[row_off:row_off + row_len],
-            grad_losses,
-            logits[row_off:row_off + row_len],
-            lse,
-            target[row_off:row_off + row_len],
-            label_smoothing,
-            logit_scale,
-            lse_square_scale,
-            softcap_val,
-            ignore_index,
-            total_classes,
-            class_start_idx,
-            n_cols,
-            logits_stride,
-            dlogits_stride,
-            grad_losses_stride,
-            row_off,
-            BLOCK_SIZE=BLOCK_SIZE,
-            HAS_SOFTCAPPING=has_softcapping,
-            num_warps=num_warps,
-        )
-def fused_cross_entropy_forward_npu(
-    logits: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float = 0.0,
-    logit_scale: float = 1.0,
-    lse_square_scale: float = 0.0,
-    logit_softcapping: float = None,
-    ignore_index: int = -100,
-    process_group=None,
-):
-    n_rows, n_cols = logits.shape
-    assert target.shape == (n_rows,)
-    world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group)
-    total_classes = world_size * n_cols
-    rank = 0 if process_group is None else torch.distributed.get_rank(process_group)
-    class_start_idx = rank * n_cols
-    if logits.stride(-1) != 1:
-        logits = logits.contiguous()
-    MAX_BLOCK_SIZE = 64 * 1024
-    BLOCK_SIZE, num_warps = _npu_block_size(n_cols, n_rows)
-    has_softcapping = logit_softcapping is not None
-    softcap_val = float(logit_softcapping) if has_softcapping else 0.0
-    n_splits = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE
-    split = world_size > 1 or n_cols > MAX_BLOCK_SIZE or n_splits > 1
-    loss_shape = (n_splits, n_rows) if n_splits > 1 else (n_rows,)
-    losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device)
-    lse = torch.empty(*loss_shape, dtype=torch.float, device=logits.device)
-    z_losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device)
-    _launch_cross_entropy_fwd(
-        losses,
-        lse,
-        z_losses,
-        logits,
-        target,
-        label_smoothing=label_smoothing,
-        logit_scale=logit_scale,
-        lse_square_scale=lse_square_scale,
-        softcap_val=softcap_val,
-        ignore_index=ignore_index,
-        total_classes=total_classes,
-        class_start_idx=class_start_idx,
-        n_cols=n_cols,
-        n_rows=n_rows,
-        logits_stride=logits.stride(0),
-        BLOCK_SIZE=BLOCK_SIZE,
-        has_softcapping=has_softcapping,
-        num_warps=num_warps,
-        split=split,
-    )
-    if split:
-        if n_splits > 1:
-            lse = torch.logsumexp(lse, dim=0)
-            losses = losses.sum(dim=0)
-        if world_size > 1:
-            lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device)
-            torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group)
-            handle_losses = torch.distributed.all_reduce(
-                losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True,
-            )
-            lse = torch.logsumexp(lse_allgather, dim=0)
-            handle_losses.wait()
-        losses += lse
-        if lse_square_scale != 0.0:
-            z_losses = lse_square_scale * lse.square()
-            z_losses.masked_fill_(target == ignore_index, 0.0)
-            losses += z_losses
-        else:
-            z_losses = torch.zeros_like(losses)
-        losses.masked_fill_(target == ignore_index, 0.0)
-    return losses, z_losses, lse, total_classes, class_start_idx
-def fused_cross_entropy_backward_npu(
-    dlogits: torch.Tensor,
-    grad_losses: torch.Tensor,
-    logits: torch.Tensor,
-    lse: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float,
-    logit_scale: float,
-    lse_square_scale: float,
-    logit_softcapping: float | None,
-    ignore_index: int,
-    total_classes: int,
-    class_start_idx: int,
-) -> torch.Tensor:
-    n_rows, n_cols = logits.shape
-    BLOCK_SIZE, num_warps = _npu_block_size(n_cols, n_rows, is_backward=True)
-    has_softcapping = logit_softcapping is not None
-    softcap_val = float(logit_softcapping) if has_softcapping else 0.0
-    _launch_cross_entropy_bwd(
-        dlogits,
-        grad_losses,
-        logits,
-        lse,
-        target,
-        label_smoothing=label_smoothing,
-        logit_scale=logit_scale,
-        lse_square_scale=lse_square_scale,
-        softcap_val=softcap_val,
-        ignore_index=ignore_index,
-        total_classes=total_classes,
-        class_start_idx=class_start_idx,
-        n_cols=n_cols,
-        n_rows=n_rows,
-        logits_stride=logits.stride(0),
-        dlogits_stride=dlogits.stride(0),
-        grad_losses_stride=grad_losses.stride(0),
-        BLOCK_SIZE=BLOCK_SIZE,
-        has_softcapping=has_softcapping,
-        num_warps=num_warps,
-    )
-    return dlogits
-class CrossEntropyLossFunctionNPU(torch.autograd.Function):
-    @staticmethod
-    @input_guard
-    def forward(
-        ctx,
-        logits,
-        target,
-        label_smoothing=0.0,
-        logit_scale=1.0,
-        lse_square_scale=0.0,
-        logit_softcapping=None,
-        ignore_index=-100,
-        inplace_backward=False,
-        process_group=None,
-    ):
-        losses, z_losses, lse, total_classes, class_start_idx = fused_cross_entropy_forward_npu(
-            logits,
-            target,
-            label_smoothing,
-            logit_scale,
-            lse_square_scale,
-            logit_softcapping,
-            ignore_index,
-            process_group,
-        )
-        ctx.save_for_backward(logits, lse, target)
-        ctx.mark_non_differentiable(z_losses)
-        ctx.label_smoothing = label_smoothing
-        ctx.logit_scale = logit_scale
-        ctx.lse_square_scale = lse_square_scale
-        ctx.logit_softcapping = logit_softcapping
-        ctx.ignore_index = ignore_index
-        ctx.total_classes = total_classes
-        ctx.class_start_idx = class_start_idx
-        ctx.inplace_backward = inplace_backward
-        return losses, z_losses
-    @staticmethod
-    @input_guard
-    def backward(ctx, grad_losses, grad_z_losses):
-        del grad_z_losses
-        logits, lse, target = ctx.saved_tensors
-        dlogits = logits if ctx.inplace_backward else torch.empty_like(logits)
-        fused_cross_entropy_backward_npu(
-            dlogits,
-            grad_losses,
-            logits,
-            lse,
-            target,
-            ctx.label_smoothing,
-            ctx.logit_scale,
-            ctx.lse_square_scale,
-            ctx.logit_softcapping,
-            ctx.ignore_index,
-            ctx.total_classes,
-            ctx.class_start_idx,
-        )
-        return dlogits, None, None, None, None, None, None, None, None, None
-def cross_entropy_loss_npu(
-    logits: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float = 0.0,
-    logit_scale: float = 1.0,
-    lse_square_scale: float = 0.0,
-    logit_softcapping: float = None,
-    ignore_index: int = -100,
-    inplace_backward: bool = False,
-    process_group=None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return CrossEntropyLossFunctionNPU.apply(
-        logits,
-        target,
-        label_smoothing,
-        logit_scale,
-        lse_square_scale,
-        logit_softcapping,
-        ignore_index,
-        inplace_backward,
-        process_group,
-    )

build/torch-cuda/modules/backends/triton_ascend/fused_kl_div.py DELETED Viewed

@@ -1,188 +0,0 @@
-# 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
-"""Fused KL divergence kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import torch.nn.functional as F
-import triton
-import triton.language as tl
-from ....ops.utils.op import exp, log
-from ....utils.ascend_ub_manager import ASCEND_MAX_GRID_DIM, compute_elementwise_block_size, compute_vocab_block_size
-_KLD_FWD_MEM_MULT = 12.0
-_ELEMENTWISE_MEM_MULT = 2.5
-STATIC_WARPS = 2
-@triton.jit
-def kl_div_kernel(
-    logits,
-    target_logits,
-    loss,
-    s_logits,
-    s_loss,
-    reduction: tl.constexpr,
-    N: tl.constexpr,
-    V: tl.constexpr,
-    BV: tl.constexpr,
-):
-    i_n = tl.program_id(0).to(tl.int64)
-    logits += i_n * s_logits
-    target_logits += i_n * s_logits
-    sm = float('-inf')
-    tm = float('-inf')
-    sd, td = 0.0, 0.0
-    NV = tl.cdiv(V, BV)
-    for iv in range(0, NV):
-        o_x = iv * BV + tl.arange(0, BV)
-        b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf'))
-        b_sm = tl.max(b_sl)
-        m_new = tl.maximum(sm, b_sm)
-        sd = sd * exp(sm - m_new) + tl.sum(exp(b_sl - m_new))
-        sm = m_new
-        b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf'))
-        b_tm = tl.max(b_tl)
-        m_new = tl.maximum(tm, b_tm)
-        td = td * exp(tm - m_new) + tl.sum(exp(b_tl - m_new))
-        tm = m_new
-    b_loss = 0.
-    for iv in range(0, NV):
-        o_x = iv * BV + tl.arange(0, BV)
-        b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf'))
-        b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf'))
-        b_sp_log = b_sl - sm - log(sd)
-        b_tp_log = b_tl - tm - log(td)
-        b_sp = exp(b_sp_log)
-        b_tp = exp(b_tp_log)
-        b_kl = tl.where(o_x < V, b_tp * (b_tp_log - b_sp_log), 0)
-        b_dl = -b_tp + b_sp
-        b_loss += tl.sum(b_kl)
-        if reduction == 'batchmean':
-            b_dl = b_dl / N
-        tl.store(logits + o_x, b_dl, mask=o_x < V)
-    if reduction == 'batchmean':
-        b_loss = b_loss / N
-    tl.store(loss + i_n * s_loss, b_loss)
-@triton.jit
-def elementwise_mul_kernel(
-    x,
-    g,
-    N: tl.constexpr,
-    B: tl.constexpr,
-):
-    i_x = tl.program_id(0).to(tl.int64)
-    o_x = i_x * B + tl.arange(0, B)
-    b_g = tl.load(g)
-    b_x = tl.load(x + o_x, mask=o_x < N)
-    tl.store(x + o_x, b_x * b_g, mask=o_x < N)
-def _npu_vocab_block_size(vocab_size: int, num_rows: int) -> int:
-    return compute_vocab_block_size(vocab_size, num_rows, _KLD_FWD_MEM_MULT)
-def fused_kl_div_forward_npu(
-    x: torch.Tensor,
-    target_x: torch.Tensor,
-    weight: torch.Tensor,
-    target_weight: torch.Tensor,
-    reduction: str = 'batchmean',
-    accumulate_grad_in_fp32: bool = True,
-):
-    device = x.device
-    N, H, V = *x.shape, weight.shape[0]
-    BV = _npu_vocab_block_size(V, N)
-    NC = min(8, triton.cdiv(V, H))
-    C = min(triton.next_power_of_2(triton.cdiv(N, NC)), ASCEND_MAX_GRID_DIM)
-    NC = triton.cdiv(N, C)
-    grad_dtype = torch.float32 if accumulate_grad_in_fp32 else weight.dtype
-    dx = torch.zeros_like(x, device=device)
-    dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None
-    loss = torch.zeros(N, dtype=torch.float32, device=device)
-    for ic in range(NC):
-        start, end = ic * C, min((ic + 1) * C, N)
-        c_sx = x[start:end]
-        c_tx = target_x[start:end]
-        c_sl = F.linear(c_sx, weight)
-        c_tl = F.linear(c_tx, target_weight)
-        if weight is not None and c_sx.dtype != grad_dtype:
-            c_sx = c_sx.to(dtype=grad_dtype)
-        c_loss = loss[start:end]
-        kl_div_kernel[(c_sx.shape[0],)](
-            logits=c_sl,
-            target_logits=c_tl,
-            loss=c_loss,
-            s_logits=c_sl.stride(-2),
-            s_loss=c_loss.stride(-1),
-            reduction=reduction,
-            N=N,
-            V=V,
-            BV=BV,
-            num_warps=STATIC_WARPS,
-        )
-        c_grad = c_sl if c_sl.is_contiguous() else c_sl.contiguous()
-        dx[start:end] = torch.mm(c_grad, weight)
-        if weight is not None:
-            grad_w = c_grad.t().to(dtype=grad_dtype)
-            grad_x = c_sx if c_sx.dtype == grad_dtype else c_sx.to(dtype=grad_dtype)
-            dw.add_(grad_w @ grad_x)
-    loss = loss.sum()
-    if dw is not None:
-        dw = dw.to(weight)
-    return loss, dx, dw
-def fused_kl_div_backward_npu(
-    do: torch.Tensor,
-    dx: torch.Tensor,
-    dw: torch.Tensor,
-):
-    if torch.ne(do, torch.tensor(1.0, device=do.device)):
-        N, H = dx.shape
-        B = compute_elementwise_block_size(N * H, _ELEMENTWISE_MEM_MULT)
-        elementwise_mul_kernel[(triton.cdiv(N * H, B),)](
-            x=dx,
-            g=do,
-            N=N*H,
-            B=B,
-            num_warps=STATIC_WARPS,
-        )
-        if dw is not None:
-            V, H = dw.shape
-            B_dw = compute_elementwise_block_size(V * H, _ELEMENTWISE_MEM_MULT)
-            elementwise_mul_kernel[(triton.cdiv(V * H, B_dw),)](
-                x=dw,
-                g=do,
-                N=V*H,
-                B=B_dw,
-                num_warps=STATIC_WARPS,
-            )
-    return dx, dw

build/torch-cuda/modules/backends/triton_ascend/fused_linear_cross_entropy.py DELETED Viewed

@@ -1,347 +0,0 @@
-# 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
-"""Fused linear cross-entropy kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import torch.nn.functional as F
-import triton
-import triton.language as tl
-from triton.language.math import tanh
-from ....ops.utils.op import exp, log
-from ....utils.ascend_ub_manager import (
-    ASCEND_MAX_GRID_DIM,
-    compute_elementwise_block_size,
-    compute_vocab_block_size,
-    iter_axis_launch_chunks,
-)
-# Fused linear CE: logsumexp forward vs gradient kernels along vocab.
-_LCE_FWD_MEM_MULT = 8.0
-_LCE_BWD_MEM_MULT = 12.0
-_ELEMENTWISE_MEM_MULT = 2.5
-STATIC_WARPS = 2
-@triton.heuristics({
-    'HAS_SCALE': lambda args: args['scale'] is not None,
-})
-@triton.jit
-def logsumexp_fwd_kernel(
-    x,
-    z,
-    scale,
-    softcapping: tl.constexpr,
-    D: tl.constexpr,
-    B: tl.constexpr,
-    ROWWISE: tl.constexpr,
-    HAS_SCALE: tl.constexpr,
-    HAS_SOFTCAPPING: tl.constexpr,
-):
-    i_n = tl.program_id(0).to(tl.int64)
-    if ROWWISE:
-        row = x + i_n * D
-        m = float('-inf')
-        d = 0.0
-        for start in range(0, D, B):
-            o = start + tl.arange(0, B)
-            b_x = tl.load(row + o, mask=o < D, other=float('-inf')).to(tl.float32)
-            if HAS_SCALE:
-                b_x = b_x * scale
-            if HAS_SOFTCAPPING:
-                b_x = softcapping * tanh(b_x / softcapping)
-            blk_max = tl.max(b_x, 0)
-            new_m = tl.maximum(m, blk_max)
-            d = d * exp(m - new_m) + tl.sum(exp(b_x - new_m), 0)
-            m = new_m
-        tl.store(z + i_n, m + log(d))
-    else:
-        i_d = tl.program_id(1).to(tl.int64)
-        o_d = i_d * B + tl.arange(0, B)
-        m_d = o_d < D
-        b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
-        if HAS_SCALE:
-            b_x = b_x * scale
-        if HAS_SOFTCAPPING:
-            b_x = softcapping * tanh(b_x / softcapping)
-        b_m = tl.max(b_x, 0)
-        b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m
-        tl.store(z + i_n * tl.cdiv(D, B) + i_d, b_z)
-@triton.jit
-def cross_entropy_kernel(
-    logits,
-    lse,
-    target,
-    loss,
-    total,
-    ignore_index,
-    label_smoothing: tl.constexpr,
-    logit_scale: tl.constexpr,
-    logit_softcapping: tl.constexpr,
-    HAS_SOFTCAPPING: tl.constexpr,
-    reduction: tl.constexpr,
-    V: tl.constexpr,
-    BV: tl.constexpr,
-):
-    i_n = tl.program_id(0).to(tl.int64)
-    NV = tl.cdiv(V, BV)
-    b_y = tl.load(target + i_n)
-    logits += i_n * V
-    if b_y == ignore_index:
-        for i in range(0, V, BV):
-            o_v = i + tl.arange(0, BV)
-            tl.store(logits + o_v, 0.0, mask=o_v < V)
-        return
-    b_l = tl.load(logits + b_y).to(tl.float32) * logit_scale
-    if HAS_SOFTCAPPING:
-        b_t_y = tanh(b_l / logit_softcapping)
-        b_l = logit_softcapping * b_t_y
-        b_softcap_deriv_y = 1.0 - b_t_y * b_t_y
-    b_lse = tl.load(lse + i_n)
-    b_loss = b_lse - b_l
-    b_z = 0.0
-    eps = label_smoothing / V
-    for iv in range(0, NV):
-        o_v = iv * BV + tl.arange(0, BV)
-        b_logits = tl.load(logits + o_v, mask=o_v < V, other=float('-inf')).to(tl.float32) * logit_scale
-        if HAS_SOFTCAPPING:
-            b_t = tanh(b_logits / logit_softcapping)
-            b_capped = logit_softcapping * b_t
-        else:
-            b_capped = b_logits
-        if label_smoothing > 0:
-            b_z += tl.sum(tl.where(o_v < V, -eps * b_capped, 0.0))
-        b_p = (exp(b_capped - b_lse) - eps) * logit_scale
-        if HAS_SOFTCAPPING:
-            b_p = b_p * (1.0 - b_t * b_t)
-        if reduction == "mean":
-            b_p = b_p / total
-        tl.store(logits + o_v, b_p, mask=o_v < V)
-    if label_smoothing > 0:
-        b_loss = b_loss * (1 - label_smoothing) + (b_z + label_smoothing * b_lse)
-    b_l = tl.load(logits + b_y)
-    if HAS_SOFTCAPPING:
-        b_sc_factor = b_softcap_deriv_y
-    else:
-        b_sc_factor = 1.0
-    if reduction == 'mean':
-        b_loss = b_loss / total
-        b_l += (label_smoothing - 1) / total * logit_scale * b_sc_factor
-    else:
-        b_l += (label_smoothing - 1) * logit_scale * b_sc_factor
-    tl.store(loss + i_n, b_loss)
-    tl.store(logits + b_y, b_l)
-@triton.jit
-def elementwise_mul_kernel(
-    x,
-    g,
-    N: tl.constexpr,
-    B: tl.constexpr,
-):
-    i_x = tl.program_id(0).to(tl.int64)
-    o_x = i_x * B + tl.arange(0, B)
-    b_g = tl.load(g)
-    b_x = tl.load(x + o_x, mask=o_x < N)
-    tl.store(x + o_x, b_x * b_g, mask=o_x < N)
-def _npu_vocab_block_size(vocab_size: int, num_rows: int, is_backward: bool = True) -> int:
-    memory_multiplier = _LCE_BWD_MEM_MULT if is_backward else _LCE_FWD_MEM_MULT
-    return compute_vocab_block_size(vocab_size, num_rows, memory_multiplier)
-def logsumexp_fwd_npu(
-    x,
-    scale: float | None = None,
-    softcapping: float | None = None,
-    dtype: torch.dtype | None = None,
-):
-    shape = x.shape
-    x = x.view(-1, shape[-1])
-    N, D = x.shape
-    B = _npu_vocab_block_size(D, N, is_backward=False)
-    has_softcapping = softcapping is not None
-    softcap_val = float(softcapping) if has_softcapping else 0.0
-    z = x.new_empty(N, dtype=torch.float)
-    for row_off, row_len in iter_axis_launch_chunks(N, 1, max_grid=ASCEND_MAX_GRID_DIM):
-        logsumexp_fwd_kernel[(row_len,)](
-            x=x[row_off:row_off + row_len],
-            z=z[row_off:row_off + row_len],
-            scale=scale,
-            softcapping=softcap_val,
-            D=D,
-            B=B,
-            ROWWISE=True,
-            HAS_SOFTCAPPING=has_softcapping,
-        )
-    z = z.view(*shape[:-1])
-    if dtype is not None and dtype != torch.float:
-        z = z.to(dtype)
-    return z
-def fused_linear_cross_entropy_forward_npu(
-    x: torch.Tensor,
-    target: torch.LongTensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor = None,
-    ignore_index: int = -100,
-    label_smoothing: float = 0.0,
-    logit_scale: float = 1.0,
-    logit_softcapping: float = None,
-    num_chunks: int = 8,
-    reduction: str = "mean",
-    use_l2warp: bool = False,
-    l2_penalty_factor: float = 1e-4,
-    accumulate_grad_in_fp32: bool = True,
-):
-    device = x.device
-    N, H, V = *x.shape, weight.shape[0]
-    BV = _npu_vocab_block_size(V, N)
-    has_softcapping = logit_softcapping is not None
-    softcap_val = float(logit_softcapping) if has_softcapping else 0.0
-    NC = min(num_chunks, triton.cdiv(V, H))
-    C = min(triton.next_power_of_2(triton.cdiv(N, NC)), ASCEND_MAX_GRID_DIM)
-    NC = triton.cdiv(N, C)
-    dx = torch.zeros_like(x, device=device)
-    grad_dtype = torch.float32 if accumulate_grad_in_fp32 else weight.dtype
-    bias_grad_dtype = None
-    if bias is not None:
-        bias_grad_dtype = torch.float32 if accumulate_grad_in_fp32 else bias.dtype
-    dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None
-    db = torch.zeros_like(bias, device=device, dtype=bias_grad_dtype) if bias is not None else None
-    loss = torch.zeros(N, device=device, dtype=torch.float)
-    total = target.ne(ignore_index).sum().item()
-    for ic in range(NC):
-        start, end = ic * C, min((ic + 1) * C, N)
-        c_x = x[start:end]
-        c_logits = F.linear(c_x, weight, bias)
-        if weight is not None and c_x.dtype != grad_dtype:
-            c_x = c_x.to(dtype=grad_dtype)
-        c_target = target[start:end]
-        c_lse = logsumexp_fwd_npu(c_logits, scale=logit_scale, softcapping=logit_softcapping, dtype=torch.float)
-        c_loss = loss[start:end]
-        if use_l2warp:
-            c_maxx, c_ids = torch.max(c_logits, -1, keepdim=True)
-        cross_entropy_kernel[(c_logits.shape[0],)](
-            logits=c_logits,
-            lse=c_lse,
-            target=c_target,
-            loss=c_loss,
-            total=total,
-            ignore_index=ignore_index,
-            label_smoothing=label_smoothing,
-            logit_scale=logit_scale,
-            logit_softcapping=softcap_val,
-            HAS_SOFTCAPPING=has_softcapping,
-            reduction=reduction,
-            V=V,
-            BV=BV,
-            num_warps=STATIC_WARPS,
-        )
-        if use_l2warp:
-            g_logits_l2 = torch.zeros_like(c_logits)
-            l2_factor = l2_penalty_factor / N
-            penalty_grad = c_maxx * l2_factor
-            g_logits_l2.scatter_(-1, c_ids, penalty_grad)
-            if weight is not None:
-                torch.addmm(
-                    input=dw,
-                    mat1=g_logits_l2.t().to(dtype=grad_dtype),
-                    mat2=c_x,
-                    out=dw,
-                )
-            if bias is not None:
-                torch.add(input=db, other=g_logits_l2.sum(0, dtype=bias_grad_dtype), out=db)
-            dx_l2_contribution = torch.mm(g_logits_l2, weight)
-        else:
-            dx_l2_contribution = 0.0
-        c_grad = c_logits if c_logits.is_contiguous() else c_logits.contiguous()
-        dx[start:end] = torch.mm(c_grad, weight) + dx_l2_contribution
-        if weight is not None:
-            grad_w = c_grad.t().to(dtype=grad_dtype)
-            grad_x = c_x if c_x.dtype == grad_dtype else c_x.to(dtype=grad_dtype)
-            dw.add_(grad_w @ grad_x)
-        if bias is not None:
-            torch.add(input=db, other=c_logits.sum(0, dtype=bias_grad_dtype), out=db)
-    loss = loss.sum()
-    if dw is not None:
-        dw = dw.to(weight)
-    if db is not None:
-        db = db.to(bias)
-    return loss, dx, dw, db
-def fused_linear_cross_entropy_backward_npu(
-    do: torch.Tensor,
-    dx: torch.Tensor,
-    dw: torch.Tensor,
-    db: torch.Tensor,
-):
-    if torch.ne(do, torch.tensor(1.0, device=do.device)):
-        N, H = dx.shape
-        B = compute_elementwise_block_size(N * H, _ELEMENTWISE_MEM_MULT)
-        elementwise_mul_kernel[(triton.cdiv(N * H, B),)](
-            x=dx,
-            g=do,
-            N=N*H,
-            B=B,
-            num_warps=STATIC_WARPS,
-        )
-        if dw is not None:
-            V, H = dw.shape
-            B_dw = compute_elementwise_block_size(V * H, _ELEMENTWISE_MEM_MULT)
-            elementwise_mul_kernel[(triton.cdiv(V * H, B_dw),)](
-                x=dw,
-                g=do,
-                N=V*H,
-                B=B_dw,
-                num_warps=STATIC_WARPS,
-            )
-        if db is not None:
-            V = db.shape[0]
-            B_db = compute_elementwise_block_size(V, _ELEMENTWISE_MEM_MULT)
-            elementwise_mul_kernel[(triton.cdiv(V, B_db),)](
-                x=db,
-                g=do,
-                N=V,
-                B=B_db,
-                num_warps=STATIC_WARPS,
-            )
-    return dx, dw, db

build/torch-cuda/modules/backends/triton_ascend/grpo.py DELETED Viewed

@@ -1,266 +0,0 @@
-# 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
-"""GRPO loss kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import triton
-import triton.language as tl
-from ....ops.utils.op import exp, log
-from ....utils import input_guard
-from ....utils.ascend_ub_manager import ASCEND_MAX_GRID_DIM, compute_vocab_block_size, iter_axis_launch_chunks
-# GRPO: use conservative multiplier covering both fwd softmax and bwd grad paths.
-_GRPO_MEM_MULT = 8.0
-STATIC_WARPS = 2
-def _npu_vocab_block_size(vocab_size: int, num_rows: int) -> int:
-    return compute_vocab_block_size(vocab_size, num_rows, _GRPO_MEM_MULT)
-@triton.jit
-def grpo_fwd_kernel(
-    logits_ptr,
-    ref_logp_ptr,
-    input_ids_ptr,
-    advantages_ptr,
-    completion_mask_ptr,
-    loss_ptr,
-    lse_ptr,
-    beta,
-    save_kl: tl.constexpr,
-    B,
-    M,
-    N,
-    L,
-    start_idx,
-    BLOCK_SIZE: tl.constexpr,
-    ROW_OFFSET: tl.constexpr,
-):
-    row_idx = tl.program_id(0) + ROW_OFFSET
-    off_b = row_idx // L
-    N = tl.cast(N, tl.int64)
-    loss_ptr += row_idx
-    completion_mask_ptr += row_idx
-    not_skip = tl.load(completion_mask_ptr).to(tl.int1)
-    if not_skip == 1:
-        ref_logp_ptr += row_idx
-        lse_ptr += row_idx
-        advantages_ptr += off_b
-        logits_ptr += N * (row_idx + off_b)
-        input_ids_ptr += row_idx + (off_b + 1) * start_idx
-        base_cols = tl.arange(0, BLOCK_SIZE)
-        m_i = -float("inf")
-        l_i = 0.0
-        for start_n in tl.range(0, N, BLOCK_SIZE):
-            cols = start_n + base_cols
-            mask = cols < N
-            logits = tl.load(logits_ptr + cols, mask=mask, other=-float('inf')).to(tl.float32)
-            m_ij = tl.max(logits)
-            new_m_i = tl.maximum(m_i, m_ij)
-            l_i = l_i * exp(m_i - new_m_i) + tl.sum(exp(logits - new_m_i))
-            m_i = new_m_i
-        lse = log(l_i) + m_i
-        idx = tl.load(input_ids_ptr)
-        x = tl.load(logits_ptr + idx).to(tl.float32)
-        advantage = tl.load(advantages_ptr).to(tl.float32)
-        ref_logp = tl.load(ref_logp_ptr)
-        logp = x - lse
-        diff = ref_logp - logp
-        kl = exp(diff) - diff - 1
-        loss = kl * beta - advantage
-        tl.store(loss_ptr, loss.to(loss_ptr.dtype.element_ty))
-        tl.store(lse_ptr, lse.to(lse_ptr.dtype.element_ty))
-        if save_kl:
-            tl.store(loss_ptr + M, kl.to(loss_ptr.dtype.element_ty))
-    else:
-        tl.store(loss_ptr, 0.0)
-        if save_kl:
-            tl.store(loss_ptr + M, 0.0)
-@triton.jit
-def grpo_bwd_kernel(
-    dloss_ptr,
-    dlogits_ptr,
-    logits_ptr,
-    ref_logp_ptr,
-    input_ids_ptr,
-    advantages_ptr,
-    completion_mask_ptr,
-    lse_ptr,
-    beta,
-    B,
-    N,
-    L,
-    start_idx,
-    BLOCK_SIZE: tl.constexpr,
-    ROW_OFFSET: tl.constexpr,
-):
-    row_idx = tl.program_id(0) + ROW_OFFSET
-    off_b = row_idx // L
-    N = tl.cast(N, tl.int64)
-    dlogits_ptr += N * (row_idx + off_b)
-    base_cols = tl.arange(0, BLOCK_SIZE)
-    completion_mask_ptr += row_idx
-    not_skip = tl.load(completion_mask_ptr).to(tl.int1)
-    if not_skip == 1:
-        lse_ptr += row_idx
-        dloss_ptr += row_idx
-        advantages_ptr += off_b
-        ref_logp_ptr += row_idx
-        logits_ptr += N * (row_idx + off_b)
-        input_ids_ptr += row_idx + (off_b + 1) * start_idx
-        dloss = tl.load(dloss_ptr).to(tl.float32)
-        lse = tl.load(lse_ptr).to(tl.float32)
-        idx = tl.load(input_ids_ptr)
-        x = tl.load(logits_ptr + idx).to(tl.float32)
-        advantage = tl.load(advantages_ptr).to(tl.float32)
-        ref_logp = tl.load(ref_logp_ptr)
-        logp = x - lse
-        dlogp = (beta * (-1.0 * exp(ref_logp - logp) + 1) - advantage) * dloss
-        for start_n in tl.range(0, N, BLOCK_SIZE):
-            cols = start_n + base_cols
-            mask = cols < N
-            logits = tl.load(logits_ptr + cols, mask=mask, other=-float('inf')).to(tl.float32)
-            probs = exp(logits - lse)
-            dlogits = tl.where(cols == idx, 1 - probs, -probs) * dlogp
-            tl.store(dlogits_ptr + cols, dlogits.to(dlogits_ptr.dtype.element_ty), mask=mask)
-    else:
-        dlogits = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
-        for start_n in tl.range(0, N, BLOCK_SIZE):
-            cols = start_n + base_cols
-            mask = cols < N
-            tl.store(dlogits_ptr + cols, dlogits.to(dlogits_ptr.dtype.element_ty), mask=mask)
-class GrpoLossNPU(torch.autograd.Function):
-    @input_guard
-    @staticmethod
-    def forward(ctx, logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl, inplace=True):
-        ctx.input_shape = logits.shape
-        B, L_ADD_1, N = ctx.input_shape
-        L = L_ADD_1 - 1
-        M = B * L
-        input_ids_start_index = input_ids.size(1) - L
-        block_size = _npu_vocab_block_size(N, M)
-        if not save_kl:
-            loss = torch.empty(B, L, device=logits.device, dtype=torch.float32)
-        else:
-            loss = torch.empty(B * 2, L, device=logits.device, dtype=torch.float32)
-        lse = torch.empty(B, L, device=logits.device, dtype=torch.float32)
-        if completion_mask is None:
-            completion_mask = torch.ones(B, L, device=logits.device, dtype=torch.int32)
-        else:
-            loss[:B].masked_fill_(completion_mask.logical_not(), 0.0)
-        for row_off, row_len in iter_axis_launch_chunks(M, 1, max_grid=ASCEND_MAX_GRID_DIM):
-            grpo_fwd_kernel[(row_len,)](
-                logits_ptr=logits,
-                ref_logp_ptr=ref_logp,
-                input_ids_ptr=input_ids,
-                advantages_ptr=advantages,
-                completion_mask_ptr=completion_mask,
-                loss_ptr=loss,
-                lse_ptr=lse,
-                beta=beta,
-                save_kl=save_kl,
-                B=B,
-                M=M,
-                N=N,
-                L=L,
-                start_idx=input_ids_start_index,
-                BLOCK_SIZE=block_size,
-                ROW_OFFSET=row_off,
-                num_warps=STATIC_WARPS,
-            )
-        ctx.beta = beta
-        ctx.save_for_backward(lse, logits, input_ids, advantages, completion_mask)
-        ctx.ref_logp = ref_logp
-        ctx.inplace = inplace
-        ctx.block_size = block_size
-        return loss
-    @input_guard
-    @staticmethod
-    def backward(ctx, dloss):
-        lse, logits, input_ids, advantages, completion_mask = ctx.saved_tensors
-        inplace = ctx.inplace
-        B, L_ADD_1, N = ctx.input_shape
-        L = L_ADD_1 - 1
-        M = B * L
-        block_size = ctx.block_size
-        input_ids_start_index = input_ids.size(1) - L
-        dlogits = logits if inplace else torch.empty_like(logits)
-        for row_off, row_len in iter_axis_launch_chunks(M, 1, max_grid=ASCEND_MAX_GRID_DIM):
-            grpo_bwd_kernel[(row_len,)](
-                dloss_ptr=dloss,
-                dlogits_ptr=dlogits,
-                logits_ptr=logits,
-                ref_logp_ptr=ctx.ref_logp,
-                input_ids_ptr=input_ids,
-                advantages_ptr=advantages,
-                completion_mask_ptr=completion_mask,
-                lse_ptr=lse,
-                beta=ctx.beta,
-                B=B,
-                N=N,
-                L=L,
-                BLOCK_SIZE=block_size,
-                start_idx=input_ids_start_index,
-                ROW_OFFSET=row_off,
-                num_warps=STATIC_WARPS,
-            )
-        dlogits[:, -1, :].fill_(0.0)
-        return dlogits.view(*ctx.input_shape), None, None, None, None, None, None, None
-def fused_grpo_loss_npu(
-    logits,
-    ref_logp,
-    input_ids,
-    advantages,
-    beta=0.1,
-    completion_mask=None,
-    save_kl=False,
-    inplace=False,
-) -> torch.Tensor:
-    out = GrpoLossNPU.apply(
-        logits,
-        ref_logp,
-        input_ids,
-        advantages,
-        beta,
-        completion_mask,
-        save_kl,
-        inplace,
-    )
-    if not save_kl:
-        return out
-    return out.chunk(2, axis=0)

build/torch-cuda/modules/backends/triton_ascend/layernorm.py DELETED Viewed

@@ -1,355 +0,0 @@
-# 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
-"""LayerNorm / RMSNorm / GroupNorm kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import triton
-import triton.language as tl
-from ....utils import get_multiprocessor_count
-from ....utils.ascend_ub_manager import ASCEND_MAX_GRID_DIM, compute_ub_block_size, iter_axis_launch_chunks
-# Peak live fp32 vectors in row-wise kernel1 (see Liger Ascend layer_norm).
-_FWD_MEM_MULT = 6.0
-_BWD_MEM_MULT = 8.0
-_UB_SAFETY_MARGIN = 0.85
-# Legacy byte cap when UB capacity cannot be detected (65536 // fp32).
-_FALLBACK_MAX_BD = 65536 // 4
-def _get_layer_norm_bd(D: int, is_forward: bool) -> int:
-    """Return power-of-2 block size for feature dim D under UB constraints."""
-    memory_multiplier = _FWD_MEM_MULT if is_forward else _BWD_MEM_MULT
-    return compute_ub_block_size(
-        D,
-        memory_multiplier,
-        safety_margin=_UB_SAFETY_MARGIN,
-        fallback=_FALLBACK_MAX_BD,
-        desired=triton.next_power_of_2(D),
-    )
-def _layer_norm_bwd_launch_config(T: int, G: int, device_index: int) -> tuple[int, int, int]:
-    """Return (NS, BS, GS) capped under Ascend grid limit."""
-    NS = min(triton.cdiv(get_multiprocessor_count(device_index), G), T // G) * G
-    NS = min(NS, ASCEND_MAX_GRID_DIM)
-    BS = triton.cdiv(T, NS) if NS > 0 else T
-    GS = NS // G if G > 0 else NS
-    return NS, BS, GS
-@triton.jit
-def layer_norm_fwd_kernel1(
-    x,
-    y,
-    w,
-    b,
-    res,
-    res_out,
-    mean,
-    rstd,
-    eps,
-    G: tl.constexpr,
-    D: tl.constexpr,
-    BD: tl.constexpr,
-    IS_RMS_NORM: tl.constexpr,
-    HAS_RESIDUAL: tl.constexpr,
-    STORE_RESIDUAL_OUT: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-):
-    i_t = tl.program_id(0)
-    i_g = i_t % G
-    x += i_t * D
-    y += i_t * D
-    if HAS_RESIDUAL:
-        res += i_t * D
-    if STORE_RESIDUAL_OUT:
-        res_out += i_t * D
-    o_d = tl.arange(0, BD)
-    m_d = o_d < D
-    b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32)
-    if HAS_RESIDUAL:
-        b_x += tl.load(res + o_d, mask=m_d, other=0.0).to(tl.float32)
-    if STORE_RESIDUAL_OUT:
-        tl.store(res_out + o_d, b_x, mask=m_d)
-    if not IS_RMS_NORM:
-        b_mean = tl.sum(b_x, axis=0) / D
-        tl.store(mean + i_t, b_mean)
-        b_xbar = tl.where(m_d, b_x - b_mean, 0.0)
-        b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
-    else:
-        b_xbar = tl.where(m_d, b_x, 0.0)
-        b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
-    b_rstd = 1 / tl.sqrt(b_var + eps)
-    tl.store(rstd + i_t, b_rstd)
-    if HAS_WEIGHT:
-        b_w = tl.load(w + i_g * D + o_d, mask=m_d).to(tl.float32)
-    if HAS_BIAS:
-        b_b = tl.load(b + i_g * D + o_d, mask=m_d).to(tl.float32)
-    b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
-    b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat
-    if HAS_BIAS:
-        b_y = b_y + b_b
-    tl.store(y + o_d, b_y, mask=m_d)
-@triton.heuristics({
-    'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None,
-})
-@triton.jit
-def layer_norm_bwd_kernel1(
-    x,
-    w,
-    b,
-    y,
-    dy,
-    dx,
-    dw,
-    db,
-    dres,
-    dres_in,
-    mean,
-    rstd,
-    T,
-    G: tl.constexpr,
-    D: tl.constexpr,
-    BS: tl.constexpr,
-    BD: tl.constexpr,
-    GS: tl.constexpr,
-    IS_RMS_NORM: tl.constexpr,
-    HAS_DRESIDUAL: tl.constexpr,
-    STORE_DRESIDUAL: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    RECOMPUTE_OUTPUT: tl.constexpr,
-):
-    i_s = tl.program_id(0)
-    i_g, i_sg = i_s // GS, i_s % GS
-    o_d = tl.arange(0, BD)
-    mask = o_d < D
-    if HAS_WEIGHT:
-        b_w = tl.load(w + i_g * D + o_d, mask=mask).to(tl.float32)
-        b_dw = tl.zeros((BD,), dtype=tl.float32)
-    if RECOMPUTE_OUTPUT and HAS_BIAS:
-        b_b = tl.load(b + i_g * D + o_d, mask=mask, other=0.0).to(tl.float32)
-    if HAS_BIAS:
-        b_db = tl.zeros((BD,), dtype=tl.float32)
-    for i_t in range(i_sg * BS * G + i_g, min((i_sg * BS + BS) * G + i_g, T), G):
-        b_x = tl.load(x + i_t * D + o_d, mask=mask, other=0).to(tl.float32)
-        b_dy = tl.load(dy + i_t * D + o_d, mask=mask, other=0).to(tl.float32)
-        if not IS_RMS_NORM:
-            b_mean = tl.load(mean + i_t)
-        b_rstd = tl.load(rstd + i_t)
-        b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
-        b_xhat = tl.where(mask, b_xhat, 0.0)
-        if RECOMPUTE_OUTPUT:
-            b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat
-            if HAS_BIAS:
-                b_y = b_y + b_b
-            tl.store(y + i_t * D + o_d, b_y, mask=mask)
-        b_wdy = b_dy
-        if HAS_WEIGHT:
-            b_wdy = b_dy * b_w
-            b_dw += b_dy * b_xhat
-        if HAS_BIAS:
-            b_db += b_dy
-        if not IS_RMS_NORM:
-            b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D
-            b_c2 = tl.sum(b_wdy, axis=0) / D
-            b_dx = (b_wdy - (b_xhat * b_c1 + b_c2)) * b_rstd
-        else:
-            b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D
-            b_dx = (b_wdy - b_xhat * b_c1) * b_rstd
-        if HAS_DRESIDUAL:
-            b_dres = tl.load(dres + i_t * D + o_d, mask=mask, other=0).to(tl.float32)
-            b_dx += b_dres
-        b_dx = tl.cast(b_dx, dtype=dx.dtype.element_ty, fp_downcast_rounding='rtne')
-        if STORE_DRESIDUAL:
-            tl.store(dres_in + i_t * D + o_d, b_dx, mask=mask)
-        tl.store(dx + i_t * D + o_d, b_dx, mask=mask)
-    if HAS_WEIGHT:
-        tl.store(dw + i_s * D + o_d, b_dw, mask=mask)
-    if HAS_BIAS:
-        tl.store(db + i_s * D + o_d, b_db, mask=mask)
-def _launch_layer_norm_fwd_kernel1(
-    x: torch.Tensor,
-    y: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    residual: torch.Tensor,
-    res_out: torch.Tensor,
-    mean: torch.Tensor,
-    rstd: torch.Tensor,
-    eps: float,
-    G: int,
-    D: int,
-    BD: int,
-    is_rms_norm: bool,
-):
-    chunk_T = x.shape[0]
-    layer_norm_fwd_kernel1[(chunk_T,)](
-        x,
-        y,
-        weight,
-        bias,
-        residual,
-        res_out,
-        mean,
-        rstd,
-        eps,
-        G=G,
-        D=D,
-        BD=BD,
-        IS_RMS_NORM=is_rms_norm,
-        HAS_RESIDUAL=residual is not None,
-        STORE_RESIDUAL_OUT=res_out is not None,
-        HAS_WEIGHT=weight is not None,
-        HAS_BIAS=bias is not None,
-    )
-def layer_norm_fwd_npu(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    eps: float = 1e-5,
-    residual: torch.Tensor = None,
-    out_dtype: torch.dtype = None,
-    residual_dtype: torch.dtype = None,
-    is_rms_norm: bool = False,
-    num_groups: int = 1,
-):
-    if residual is not None:
-        residual_dtype = residual.dtype
-    T, D, G = *x.shape, num_groups
-    if residual is not None:
-        assert residual.shape == (T, D)
-    if weight is not None:
-        assert weight.shape == (G * D,)
-    if bias is not None:
-        assert bias.shape == (G * D,)
-    y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
-    if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype):
-        res_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
-    else:
-        res_out = None
-    mean = torch.empty((T,), dtype=torch.float, device=x.device) if not is_rms_norm else None
-    rstd = torch.empty((T,), dtype=torch.float, device=x.device)
-    BD = _get_layer_norm_bd(D, is_forward=True)
-    if D > BD:
-        raise RuntimeError(
-            f"LayerNorm feature dim {D} exceeds UB-safe block size {BD}. "
-            "Column-tiled kernels are not yet implemented for this size."
-        )
-    # Ascend: use row-wise kernel1 (no make_block_ptr) for all feature dims.
-    # Split along rows when T exceeds the Ascend grid limit.
-    for row_start, row_len in iter_axis_launch_chunks(T, 1, max_grid=ASCEND_MAX_GRID_DIM):
-        row_end = row_start + row_len
-        _launch_layer_norm_fwd_kernel1(
-            x[row_start:row_end],
-            y[row_start:row_end],
-            weight,
-            bias,
-            None if residual is None else residual[row_start:row_end],
-            None if res_out is None else res_out[row_start:row_end],
-            None if mean is None else mean[row_start:row_end],
-            rstd[row_start:row_end],
-            eps,
-            G,
-            D,
-            BD,
-            is_rms_norm,
-        )
-    return y, mean, rstd, res_out if res_out is not None else x
-def layer_norm_bwd_npu(
-    dy: torch.Tensor,
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    mean: torch.Tensor = None,
-    rstd: torch.Tensor = None,
-    dres: torch.Tensor = None,
-    has_residual: bool = False,
-    is_rms_norm: bool = False,
-    x_dtype: torch.dtype = None,
-    recompute_output: bool = False,
-    num_groups: int = 1,
-):
-    T, D, G = *x.shape, num_groups
-    assert dy.shape == (T, D)
-    if dres is not None:
-        assert dres.shape == (T, D)
-    if weight is not None:
-        assert weight.shape == (G * D,)
-    if bias is not None:
-        assert bias.shape == (G * D,)
-    dx = torch.empty_like(x) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
-    dres_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None
-    y = torch.empty(T, D, dtype=dy.dtype, device=dy.device) if recompute_output else None
-    BD = _get_layer_norm_bd(D, is_forward=False)
-    if D > BD:
-        raise RuntimeError(
-            f"LayerNorm feature dim {D} exceeds UB-safe block size {BD}. "
-            "Column-tiled kernels are not yet implemented for this size."
-        )
-    NS, BS, GS = _layer_norm_bwd_launch_config(T, G, x.device.index)
-    dw = torch.empty((NS, D), dtype=torch.float, device=weight.device) if weight is not None else None
-    db = torch.empty((NS, D), dtype=torch.float, device=bias.device) if bias is not None else None
-    grid = (NS,)
-    layer_norm_bwd_kernel1[grid](
-        x,
-        weight,
-        bias,
-        y,
-        dy,
-        dx,
-        dw,
-        db,
-        dres,
-        dres_in,
-        mean,
-        rstd,
-        T=T,
-        G=G,
-        D=D,
-        BS=BS,
-        BD=BD,
-        GS=GS,
-        IS_RMS_NORM=is_rms_norm,
-        HAS_DRESIDUAL=dres is not None,
-        STORE_DRESIDUAL=dres_in is not None,
-        HAS_WEIGHT=weight is not None,
-        HAS_BIAS=bias is not None,
-    )
-    dw = dw.view(G, -1, D).sum(1).to(weight).view_as(weight) if weight is not None else None
-    db = db.view(G, -1, D).sum(1).to(bias).view_as(bias) if bias is not None else None
-    if has_residual and dx.dtype == x.dtype:
-        dres_in = dx
-    return (dx, dw, db, dres_in) if not recompute_output else (dx, dw, db, dres_in, y)

build/torch-cuda/modules/backends/triton_ascend/rotary.py DELETED Viewed

@@ -1,211 +0,0 @@
-# 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
-"""Rotary embedding kernels adapted for triton-ascend on Huawei NPU."""
-import torch
-import triton
-import triton.language as tl
-from ....ops.utils import prepare_chunk_indices
-from ....utils import autotune_cache_kwargs, get_multiprocessor_count
-from ....utils.ascend_ub_manager import (
-    ASCEND_MAX_GRID_DIM,
-    compute_grid_limited_tile_size,
-    compute_row_tile_block_size,
-    max_grid_axis_chunks,
-)
-# Peak live fp32 tiles in rotary kernel: cos, sin, x0, x1, o0, o1.
-_ROTARY_MEM_MULT = 6.0
-_ROTARY_SAFETY_MARGIN = 0.90
-# Ascend vector UB is small; large num_warps / stages explodes compile-time UB (see bishengir ub overflow).
-NUM_WARPS_AUTOTUNE = [2, 4]
-NUM_STAGES_AUTOTUNE = [1, 2]
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
-        for num_warps in NUM_WARPS_AUTOTUNE
-        for num_stages in NUM_STAGES_AUTOTUNE
-    ],
-    key=['B', 'H', 'D', 'INTERLEAVED'],
-    **autotune_cache_kwargs,
-)
-@triton.jit(do_not_specialize=['T'])
-def rotary_embedding_kernel(
-    x,
-    cos,
-    sin,
-    y,
-    cu_seqlens,
-    chunk_indices,
-    seq_offsets,
-    T,
-    B: tl.constexpr,
-    H: tl.constexpr,
-    D: tl.constexpr,
-    R: tl.constexpr,
-    TR: tl.constexpr,
-    BT: tl.constexpr,
-    BD: tl.constexpr,
-    IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr,
-    IS_VARLEN: tl.constexpr,
-    INTERLEAVED: tl.constexpr,
-    CONJUGATE: tl.constexpr,
-    NT_OFFSET: tl.constexpr,
-):
-    i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    i_t += NT_OFFSET
-    if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
-        bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1)
-        T = eos - bos
-        x = x + bos * H*D + i_h * D
-        y = y + bos * H*D + i_h * D
-    else:
-        i_n = i_b
-        x = x + i_n * T*H*D + i_h * D
-        y = y + i_n * T*H*D + i_h * D
-    if i_t * BT >= T:
-        return
-    o_t = i_t * BT + tl.arange(0, BT)
-    if not IS_SEQLEN_OFFSETS_TENSOR:
-        o_cs = o_t + seq_offsets
-    else:
-        o_cs = o_t + tl.load(seq_offsets + i_n)
-    m_t = (o_t >= 0) & (o_t < T) & (o_cs >= 0) & (o_cs < TR)
-    if not INTERLEAVED:
-        o_r = tl.arange(0, BD // 2)
-        p_x = x + o_t[:, None] * H*D + o_r[None, :]
-        p_cos = cos + (o_cs[:, None] * R + o_r[None, :])
-        p_sin = sin + (o_cs[:, None] * R + o_r[None, :])
-        mask = m_t[:, None] & (o_r < R)[None, :]
-        b_cos = tl.load(p_cos, mask=mask, other=1.0).to(tl.float32)
-        b_sin = tl.load(p_sin, mask=mask, other=0.0).to(tl.float32)
-        b_x0 = tl.load(p_x, mask=mask, other=0.0).to(tl.float32)
-        b_x1 = tl.load(p_x + R, mask=mask, other=0.0).to(tl.float32)
-        if CONJUGATE:
-            b_sin = -b_sin
-        b_o0 = b_x0 * b_cos - b_x1 * b_sin
-        b_o1 = b_x0 * b_sin + b_x1 * b_cos
-        p_y = y + (o_t[:, None] * H*D + o_r[None, :])
-        tl.store(p_y, b_o0, mask=mask)
-        tl.store(p_y + R, b_o1, mask=mask)
-    else:
-        o_d = tl.arange(0, BD)
-        o_d_swap = o_d + ((o_d + 1) % 2) * 2 - 1
-        o_d_repeat = tl.arange(0, BD) // 2
-        p_x0 = x + o_t[:, None] * H*D + o_d[None, :]
-        p_x1 = x + o_t[:, None] * H*D + o_d_swap[None, :]
-        p_cos = cos + (o_cs[:, None] * R + o_d_repeat[None, :])
-        p_sin = sin + (o_cs[:, None] * R + o_d_repeat[None, :])
-        mask = m_t[:, None] & (o_d_repeat < R)[None, :]
-        b_cos = tl.load(p_cos, mask=mask, other=1.0).to(tl.float32)
-        b_sin = tl.load(p_sin, mask=mask, other=0.0).to(tl.float32)
-        b_x0 = tl.load(p_x0, mask=mask, other=0.0).to(tl.float32)
-        b_x1 = tl.load(p_x1, mask=mask, other=0.0).to(tl.float32)
-        if CONJUGATE:
-            b_sin = -b_sin
-        b_o0 = b_x0 * b_cos
-        b_o1 = b_x1 * b_sin
-        b_y = tl.where(o_d[None, :] % 2 == 0, b_o0 - b_o1, b_o0 + b_o1)
-        p_y = y + (o_t[:, None] * H*D + o_d[None, :])
-        tl.store(p_y, b_y, mask=mask)
-def rotary_embedding_fwdbwd_npu(
-    x: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    seqlen_offsets: int | torch.Tensor = 0,
-    cu_seqlens: torch.Tensor | None = None,
-    interleaved: bool = False,
-    inplace: bool = False,
-    conjugate: bool = False,
-    chunk_indices: torch.LongTensor | None = None,
-) -> torch.Tensor:
-    is_varlen = cu_seqlens is not None
-    B, T, H, D = x.shape
-    N = B if not is_varlen else cu_seqlens.shape[0] - 1
-    TR, R = cos.shape
-    R2 = R * 2
-    assert D <= 256, "Only support D <= 256"
-    assert TR >= T, f"TR must be >= T, got {TR} and {T}"
-    assert cos.dtype == sin.dtype, f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}"
-    assert x.dtype == cos.dtype, f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}"
-    if isinstance(seqlen_offsets, torch.Tensor):
-        assert seqlen_offsets.shape == (N,)
-        assert seqlen_offsets.dtype in [torch.int32, torch.int64]
-    else:
-        assert seqlen_offsets + T <= TR
-    y = torch.zeros_like(x) if not inplace else x
-    if R2 < D and not inplace:
-        y[..., R2:].copy_(x[..., R2:])
-    BD = triton.next_power_of_2(R2)
-    desired_bt = triton.next_power_of_2(triton.cdiv(T, get_multiprocessor_count(x.device.index)))
-    bt_cap = compute_row_tile_block_size(
-        desired_bt,
-        R2,
-        _ROTARY_MEM_MULT,
-        safety_margin=_ROTARY_SAFETY_MARGIN,
-        dtype_size=x.element_size(),
-        fallback=16 if R >= 128 else (32 if R >= 64 else 64),
-        min_block=1,
-    )
-    BT = min(bt_cap, desired_bt)
-    BT = compute_grid_limited_tile_size(T, B * H, BT, max_grid=ASCEND_MAX_GRID_DIM)
-    if chunk_indices is None and is_varlen:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
-    NT = len(chunk_indices) if is_varlen else triton.cdiv(T, BT)
-    kernel_kwargs = dict(
-        x=x,
-        cos=cos,
-        sin=sin,
-        y=y,
-        cu_seqlens=cu_seqlens,
-        chunk_indices=chunk_indices,
-        seq_offsets=seqlen_offsets,
-        B=B,
-        T=T,
-        H=H,
-        D=D,
-        R=R,
-        TR=TR,
-        BT=BT,
-        BD=BD,
-        IS_SEQLEN_OFFSETS_TENSOR=isinstance(seqlen_offsets, torch.Tensor),
-        IS_VARLEN=is_varlen,
-        INTERLEAVED=interleaved,
-        CONJUGATE=conjugate,
-    )
-    max_nt = max_grid_axis_chunks(NT, B * H, max_grid=ASCEND_MAX_GRID_DIM)
-    for nt_off in range(0, NT, max_nt):
-        nt_len = min(max_nt, NT - nt_off)
-        if is_varlen:
-            kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len]
-            kernel_kwargs['NT_OFFSET'] = 0
-        else:
-            kernel_kwargs['chunk_indices'] = chunk_indices
-            kernel_kwargs['NT_OFFSET'] = nt_off
-        rotary_embedding_kernel[(nt_len, B, H)](**kernel_kwargs)
-    return y

build/{torch-xpu/modules/conv/triton → torch-cuda/modules/causal_conv1d}/__init__.py RENAMED Viewed

@@ -5,20 +5,27 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-from .ops import (
     CausalConv1dFunction,
     causal_conv1d_bwd,
     causal_conv1d_fwd,
     causal_conv1d_update,
     causal_conv1d_update_states,
-    compute_dh0_triton,
 )
 __all__ = [
     'CausalConv1dFunction',
     'causal_conv1d_bwd',
     'causal_conv1d_fwd',
     'causal_conv1d_update',
     'causal_conv1d_update_states',
-    'compute_dh0_triton',
 ]

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+from ...backends import GluonBackend
+from ...modules.causal_conv1d.ops import (
     CausalConv1dFunction,
+    ShortConvolution,
+    causal_conv1d,
     causal_conv1d_bwd,
     causal_conv1d_fwd,
     causal_conv1d_update,
     causal_conv1d_update_states,
 )
 __all__ = [
     'CausalConv1dFunction',
+    'ShortConvolution',
+    'causal_conv1d',
     'causal_conv1d_bwd',
     'causal_conv1d_fwd',
     'causal_conv1d_update',
     'causal_conv1d_update_states',
 ]
+if GluonBackend.is_available():
+    from ...modules.causal_conv1d.backends import gluon  # noqa: F401

build/torch-cuda/modules/causal_conv1d/backends/__init__.py ADDED Viewed

	@@ -0,0 +1,6 @@

+# 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

build/torch-cuda/modules/{conv → causal_conv1d/backends}/cuda/__init__.py RENAMED Viewed

@@ -5,10 +5,6 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-from .ops import FastCausalConv1dFn, causal_conv1d_cuda, fast_causal_conv1d_fn
-__all__ = [
-    'FastCausalConv1dFn',
-    'causal_conv1d_cuda',
-    'fast_causal_conv1d_fn',
-]

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+from .....modules.causal_conv1d.backends.cuda.ops import FastCausalConv1dFn, causal_conv1d_cuda, fast_causal_conv1d_fn
+__all__ = ['FastCausalConv1dFn', 'causal_conv1d_cuda', 'fast_causal_conv1d_fn']

build/torch-cuda/modules/causal_conv1d/backends/cuda/cp.py ADDED Viewed

	@@ -0,0 +1,98 @@

+# 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 __future__ import annotations
+from typing import TYPE_CHECKING
+import torch
+import torch.nn.functional as F
+from .....ops.cp import conv_cp_send_recv_bwd, conv_cp_send_recv_fwd
+from .....ops.utils import prepare_sequence_ids
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+    from .....ops.cp import FLACPContext
+class ConvBoundaryExchangeFunction(torch.autograd.Function):
+    """Exchange convolution history and return its gradients to the sending rank."""
+    @staticmethod
+    def forward(ctx, tails: torch.Tensor, group: ProcessGroup) -> torch.Tensor:
+        ctx.group = group
+        return conv_cp_send_recv_fwd(tails=tails, group=group)
+    @staticmethod
+    def backward(ctx, dheads: torch.Tensor) -> tuple[torch.Tensor, None]:
+        return conv_cp_send_recv_bwd(d_initial_state=dheads.contiguous(), group=ctx.group), None
+def causal_conv1d_cuda_cp(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    residual: torch.Tensor | None,
+    activation: str | None,
+    cp_context: FLACPContext,
+) -> torch.Tensor:
+    """Apply CUDA convolution to contiguous CP inputs of shape `[1, T, D]`.
+    Each rank needs at least `W - 1` tokens, with `W` in {2, 3, 4}.
+    Packed sequence boundaries mask both history values and their gradients.
+    """
+    from .....modules.causal_conv1d.backends.cuda.ops import causal_conv1d_fn_cuda
+    if causal_conv1d_fn_cuda is None:
+        raise ImportError("CUDA CP requires the causal-conv1d package")
+    if cp_context.layout != 'contiguous':
+        raise NotImplementedError("CUDA convolution supports only contiguous CP")
+    width = weight.shape[-1]
+    if width not in (2, 3, 4):
+        raise NotImplementedError("CUDA CP convolution supports kernel widths 2, 3, and 4")
+    if x.ndim != 3 or x.shape[0] != 1:
+        raise ValueError("CUDA CP convolution requires input shape [1, T, D]")
+    if x.shape[1] < width - 1:
+        raise ValueError("CUDA CP convolution requires at least kernel_width - 1 tokens per rank")
+    if cp_context.conv1d_kernel_size != width:
+        raise ValueError("CP context kernel width must match the convolution weight")
+    if activation not in (None, 'silu', 'swish'):
+        raise NotImplementedError("activation must be None, silu, or swish")
+    initial_state = None
+    if cp_context.group is not None:
+        tails = x[0, -(width - 1):].contiguous()
+        heads = ConvBoundaryExchangeFunction.apply(tails, cp_context.group)
+        valid_len = min(width - 1, cp_context.pre_num_conv_tokens or 0)
+        # history must not cross the first local document's global start
+        mask = torch.arange(width - 1, device=x.device) >= width - 1 - valid_len
+        initial_state = heads.masked_fill(~mask[:, None], 0).unsqueeze(0)
+    channels = x.shape[-1]
+    # the CUDA channel-last kernels require strides divisible by eight
+    padding = -channels % 8
+    x = F.pad(x, (0, padding)) if padding else x.contiguous()
+    if padding:
+        weight = F.pad(weight, (0, 0, 0, padding))
+        bias = F.pad(bias, (0, padding)) if bias is not None else None
+        initial_state = F.pad(initial_state, (0, padding)) if initial_state is not None else None
+    seq_idx = prepare_sequence_ids(cu_seqlens=cp_context.cu_seqlens, cu_seqlens_cpu=cp_context.cu_seqlens_cpu)
+    seq_idx = seq_idx.to(torch.int32).unsqueeze(0)
+    y = causal_conv1d_fn_cuda(
+        x=x.transpose(1, 2),
+        weight=weight,
+        bias=bias,
+        seq_idx=seq_idx,
+        initial_states=initial_state.transpose(1, 2) if initial_state is not None else None,
+        activation=activation,
+    )
+    y = y.transpose(1, 2)[..., :channels]
+    return y + residual if residual is not None else y

build/{torch-xpu/modules/conv → torch-cuda/modules/causal_conv1d/backends}/cuda/ops.py RENAMED Viewed

@@ -5,14 +5,14 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-"""CUDA-based mixed-mode implementation for causal convolution."""
 import torch
 from einops import rearrange
-from ....modules.conv.triton import causal_conv1d_update_states
-from ....ops.utils import prepare_sequence_ids
-from ....utils import input_guard
 try:
     from causal_conv1d.cpp_functions import causal_conv1d_bwd_function
@@ -26,45 +26,27 @@ except ImportError:
 class FastCausalConv1dFn(torch.autograd.Function):
-    """
-    Mixed-mode (Mix) Causal Convolution Implementation - Combining Triton Forward and CUDA Backward Propagation
-    This class implements forward propagation using FLA's Triton kernel, while using the optimized
-    implementation from TriDao's causal_conv1d CUDA package for backward propagation.
-    This hybrid strategy combines the advantages of both technologies:
-    - Forward: Uses FLA's Triton implementation, optimized for the FLA framework
-    - Backward: Uses TriDao's causal_conv1d_bwd_function CUDA implementation for faster speed
-    Performance Benefits:
-    - CUDA backward implementation is typically faster than the Triton version, reducing training time
-    - Maintains the flexibility and compatibility of forward propagation
-    Note:
-    - Input/Output format is (batch, seqlen, dim)
-    - Backward propagation requires causal_conv1d package: pip install causal-conv1d
-    - Supports SILU/Swish activation functions
-    - Current limitations (not yet supported):
-        * output_final_state must be False
-        * initial_states must be None
-        * residual must be None
     """
     @staticmethod
     @input_guard(no_guard_contiguous=["x"])
     def forward(
         ctx,
-        x,
-        weight,
-        bias=None,
         residual: torch.Tensor | None = None,
-        initial_states=None,
-        output_final_state=False,
-        activation=None,
         cu_seqlens: torch.LongTensor | None = None,
         cu_seqlens_cpu: torch.LongTensor | None = None,
         chunk_indices: torch.LongTensor | None = None,
         seq_idx: torch.LongTensor | None = None,
-    ):
         if activation not in [None, "silu", "swish"]:
             raise NotImplementedError("activation must be None, silu, or swish")
         assert output_final_state is False, "output_final_state must be False for FastCausalConv1dFn"
@@ -73,12 +55,11 @@ class FastCausalConv1dFn(torch.autograd.Function):
         bias = bias.contiguous() if bias is not None else None
         if cu_seqlens is not None and seq_idx is None:
-            seq_idx = prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(
-                torch.int32).unsqueeze(0)
         seq_idx = seq_idx.contiguous() if seq_idx is not None else None
-        # Import here to avoid circular dependency
-        from ....modules.conv.triton.ops import causal_conv1d_fwd
         ctx.activation = activation in ["silu", "swish"]
         out, _ = causal_conv1d_fwd(
@@ -96,14 +77,12 @@ class FastCausalConv1dFn(torch.autograd.Function):
         ctx.save_for_backward(x, weight, bias, seq_idx, initial_states)
         ctx.return_final_states = output_final_state
-        ctx.return_dinitial_states = (
-            initial_states is not None and initial_states.requires_grad
-        )
         return out, None
     @staticmethod
     @input_guard
-    def backward(ctx, dout, *args):
         x, weight, bias, seq_idx, initial_states = ctx.saved_tensors
         dx = torch.empty_like(x, memory_format=torch.contiguous_format)
         x = rearrange(x, 'b t d -> b d t')
@@ -113,35 +92,20 @@ class FastCausalConv1dFn(torch.autograd.Function):
         if dout.stride(2) != 1 and dout.stride(1) != 1:
             dout = dout.contiguous()
-        # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the
-        # backward of conv1d with the backward of chunk).
-        # Here we just pass in None and dx will be allocated in the C++ code.
         dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function(
-            x,
-            weight,
-            bias,
-            dout,
-            seq_idx,
-            initial_states,
-            dfinal_states,
-            dx,
-            ctx.return_dinitial_states,
-            ctx.activation,
         )
         dx = rearrange(dx, 'b d t -> b t d')
-        return (
-            dx,
-            dweight,
-            dbias if bias is not None else None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
 def fast_causal_conv1d_fn(
@@ -156,17 +120,36 @@ def fast_causal_conv1d_fn(
     cu_seqlens_cpu: torch.LongTensor | None = None,
     chunk_indices: torch.LongTensor | None = None,
     seq_idx: torch.LongTensor | None = None,
-):
-    """
-    x: (batch, seqlen, dim)
-    weight: (dim, width)
-    bias: (dim,)
-    seq_idx: (batch, seqlen)
-    initial_states: (batch, dim, width - 1)
-    final_states_out: (batch, dim, width - 1), to be written to
-    activation: either None or "silu" or "swish"
-    out: (batch, seqlen, dim)
     """
     assert causal_conv1d_bwd_function is not None, "causal_conv1d_bwd_function is not available"
     return FastCausalConv1dFn.apply(
@@ -195,7 +178,7 @@ def causal_conv1d_cuda(
     cu_seqlens: torch.Tensor | None = None,
     cu_seqlens_cpu: torch.LongTensor | None = None,
     **kwargs,
-):
     assert causal_conv1d_fn_cuda is not None, "causal_conv1d_fn_cuda is not available"
     seq_idx = kwargs.get('seq_idx')
     if cu_seqlens is not None or seq_idx is not None:
@@ -205,29 +188,27 @@ def causal_conv1d_cuda(
         x = x.contiguous()
     x_conv1d = rearrange(x, 'b t d -> b d t')
     if cu_seqlens is not None and seq_idx is None:
-        seq_idx = prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(torch.int32).unsqueeze(0)
     y = causal_conv1d_fn_cuda(
         x=x_conv1d,
         weight=weight,
         bias=bias,
-        activation=activation,
         seq_idx=seq_idx,
         initial_states=None,
         return_final_states=False,
     )
     y = rearrange(y, 'b d t -> b t d')
     if output_final_state:
-        final_state = causal_conv1d_update_states(
-            x=x,
-            state_len=W,
-            initial_state=initial_state,
-            cu_seqlens=cu_seqlens,
-        )
     else:
         final_state = None
     if residual is not None:
         y.add_(residual)
     return y, final_state

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+"""CUDA and mixed Triton/CUDA implementations of causal convolution."""
 import torch
 from einops import rearrange
+from .....modules.causal_conv1d import causal_conv1d_update_states
+from .....ops.utils import prepare_sequence_ids
+from .....utils import input_guard
 try:
     from causal_conv1d.cpp_functions import causal_conv1d_bwd_function
 class FastCausalConv1dFn(torch.autograd.Function):
+    """Causal convolution with Triton forward and CUDA backward on `[B, T, D]` inputs.
+    Backward requires the `causal-conv1d` package. Activation may be `None`, `silu`, or `swish`.
+    Residuals and initial/final states are unsupported.
     """
     @staticmethod
     @input_guard(no_guard_contiguous=["x"])
     def forward(
         ctx,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        bias: torch.Tensor | None = None,
         residual: torch.Tensor | None = None,
+        initial_states: torch.Tensor | None = None,
+        output_final_state: bool = False,
+        activation: str | None = None,
         cu_seqlens: torch.LongTensor | None = None,
         cu_seqlens_cpu: torch.LongTensor | None = None,
         chunk_indices: torch.LongTensor | None = None,
         seq_idx: torch.LongTensor | None = None,
+    ) -> tuple[torch.Tensor, None]:
         if activation not in [None, "silu", "swish"]:
             raise NotImplementedError("activation must be None, silu, or swish")
         assert output_final_state is False, "output_final_state must be False for FastCausalConv1dFn"
         bias = bias.contiguous() if bias is not None else None
         if cu_seqlens is not None and seq_idx is None:
+            seq_idx = prepare_sequence_ids(cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(torch.int32).unsqueeze(0)
         seq_idx = seq_idx.contiguous() if seq_idx is not None else None
+        # import here to avoid circular dependency.
+        from .....modules.causal_conv1d import causal_conv1d_fwd
         ctx.activation = activation in ["silu", "swish"]
         out, _ = causal_conv1d_fwd(
         ctx.save_for_backward(x, weight, bias, seq_idx, initial_states)
         ctx.return_final_states = output_final_state
+        ctx.return_dinitial_states = initial_states is not None and initial_states.requires_grad
         return out, None
     @staticmethod
     @input_guard
+    def backward(ctx, dout: torch.Tensor, *args) -> tuple[torch.Tensor | None, ...]:
         x, weight, bias, seq_idx, initial_states = ctx.saved_tensors
         dx = torch.empty_like(x, memory_format=torch.contiguous_format)
         x = rearrange(x, 'b t d -> b d t')
         if dout.stride(2) != 1 and dout.stride(1) != 1:
             dout = dout.contiguous()
         dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function(
+            x=x,
+            weight=weight,
+            bias=bias,
+            dout=dout,
+            seq_idx=seq_idx,
+            initial_states=initial_states,
+            dfinal_states=dfinal_states,
+            dx=dx,
+            return_dinitial_states=ctx.return_dinitial_states,
+            silu_activation=ctx.activation,
         )
         dx = rearrange(dx, 'b d t -> b t d')
+        return (dx, dweight, dbias if bias is not None else None, None, None, None, None, None, None, None, None)
 def fast_causal_conv1d_fn(
     cu_seqlens_cpu: torch.LongTensor | None = None,
     chunk_indices: torch.LongTensor | None = None,
     seq_idx: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply mixed Triton/CUDA causal convolution to `[B, T, D]` inputs.
+    Args:
+        x (torch.Tensor):
+            Input of shape `[B, T, D]`.
+        weight (torch.Tensor, Optional):
+            Convolution weights of shape `[D, W]`. Default: `None`.
+        bias (torch.Tensor, Optional):
+            Bias of shape `[D]`. Default: `None`.
+        residual (torch.Tensor, Optional):
+            Unsupported by this backend; must be `None`. Default: `None`.
+        initial_state (torch.Tensor, Optional):
+            Unsupported by this backend; must be `None`. Default: `None`.
+        output_final_state (bool, Optional):
+            Unsupported by this backend; must be `False`. Default: `False`.
+        activation (str, Optional):
+            Activation applied to the output: `None`, `silu`, or `swish`. Default: `None`.
+        cu_seqlens (torch.Tensor, Optional):
+            Cumulative sequence lengths for packed inputs. Default: `None`.
+        cu_seqlens_cpu (torch.LongTensor, Optional):
+            CPU copy of `cu_seqlens`. Default: `None`.
+        chunk_indices (torch.LongTensor, Optional):
+            Precomputed sequence chunk indices. Default: `None`.
+        seq_idx (torch.LongTensor, Optional):
+            Sequence IDs of shape `[B, T]`. Derived from `cu_seqlens` when omitted. Default: `None`.
+    Returns:
+        tuple[torch.Tensor, None]:
+            Output of shape `[B, T, D]` and no final state.
     """
     assert causal_conv1d_bwd_function is not None, "causal_conv1d_bwd_function is not available"
     return FastCausalConv1dFn.apply(
     cu_seqlens: torch.Tensor | None = None,
     cu_seqlens_cpu: torch.LongTensor | None = None,
     **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
     assert causal_conv1d_fn_cuda is not None, "causal_conv1d_fn_cuda is not available"
     seq_idx = kwargs.get('seq_idx')
     if cu_seqlens is not None or seq_idx is not None:
         x = x.contiguous()
     x_conv1d = rearrange(x, 'b t d -> b d t')
     if cu_seqlens is not None and seq_idx is None:
+        seq_idx = prepare_sequence_ids(cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(torch.int32).unsqueeze(0)
     y = causal_conv1d_fn_cuda(
         x=x_conv1d,
         weight=weight,
         bias=bias,
         seq_idx=seq_idx,
         initial_states=None,
         return_final_states=False,
+        activation=activation,
     )
     y = rearrange(y, 'b d t -> b t d')
     if output_final_state:
+        final_state = causal_conv1d_update_states(x=x, state_len=W, initial_state=initial_state, cu_seqlens=cu_seqlens)
     else:
         final_state = None
     if residual is not None:
         y.add_(residual)
     return y, final_state
+__all__ = ['FastCausalConv1dFn', 'causal_conv1d_cuda', 'fast_causal_conv1d_fn']

build/torch-cuda/modules/causal_conv1d/backends/gluon.py ADDED Viewed

	@@ -0,0 +1,417 @@

+# 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
+"""Causal depthwise convolution with register reuse and fused backward recomputation."""
+import torch
+import triton
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+from ....backends import GluonBackend, register
+from ....modules.causal_conv1d import causal_conv1d_bwd, causal_conv1d_fwd
+from ....ops.utils import prepare_chunk_indices
+from ....utils import input_guard
+@gluon.jit(do_not_specialize=['T'])
+def causal_conv1d_fwd_kernel(
+    x,
+    y,
+    weight,
+    bias,
+    residual,
+    cu_seqlens,
+    initial_state,
+    chunk_indices,
+    T,
+    SXN: gl.constexpr,
+    SXT: gl.constexpr,
+    D: gl.constexpr,
+    W: gl.constexpr,
+    BT: gl.constexpr,
+    BD: gl.constexpr,
+    BC: gl.constexpr,
+    ACTIVATION: gl.constexpr,
+    NUM_SPLITS: gl.constexpr = 1,
+):
+    layout: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, BD // 32],
+        threads_per_warp=[1, 32],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    i_d = gl.program_id(0).to(gl.int64)
+    i_t = gl.program_id(1).to(gl.int64)
+    i_b = gl.program_id(2).to(gl.int64)
+    i_split = i_t % NUM_SPLITS
+    i_t //= NUM_SPLITS
+    if cu_seqlens is not None:
+        i_n = gl.load(chunk_indices + 2 * i_t).to(gl.int64)
+        i_t = gl.load(chunk_indices + 2 * i_t + 1).to(gl.int64)
+        bos = gl.load(cu_seqlens + i_n).to(gl.int64)
+        T = gl.load(cu_seqlens + i_n + 1).to(gl.int64) - bos
+        p_x = x + bos * SXT
+    else:
+        i_n = i_b
+        bos = i_b * T
+        p_x = x + i_b * SXN
+    o_d = i_d * BD + gl.arange(0, BD, layout=gl.SliceLayout(0, layout)).to(gl.int64)
+    o_t = (i_t * NUM_SPLITS + i_split) * BT
+    o_t += gl.arange(0, BT // BC, layout=gl.SliceLayout(1, layout)).to(gl.int64) * BC
+    b_x = ()
+    b_w = ()
+    for i_w in gl.static_range(W):
+        b_w += (gl.load(weight + o_d * W + i_w, mask=o_d < D, other=0).to(gl.float32),)
+    # each row unrolls BC outputs from one register window; adjacent outputs reuse inputs.
+    # with W=4:
+    # y[t]   <- x[t-3] x[t-2] x[t-1] x[t]
+    # y[t+1] <-        x[t-2] x[t-1] x[t] x[t+1]
+    for i_x in gl.static_range(BC + W - 1):
+        o_x = o_t + i_x - W + 1
+        b_xi = gl.load(
+            pointer=p_x + o_x[:, None] * SXT + o_d[None, :],
+            mask=(o_x[:, None] >= 0) & (o_x[:, None] < T) & (o_d[None, :] < D),
+            other=0,
+        ).to(gl.float32)
+        if initial_state is not None:
+            b_xi += gl.load(
+                pointer=initial_state + i_n * D * W + o_d[None, :] * W + (o_x[:, None] + W),
+                mask=(o_x[:, None] < 0) & (o_x[:, None] >= 1 - W) & (o_d[None, :] < D),
+                other=0,
+            ).to(gl.float32)
+        b_x += (b_xi,)
+    if bias is not None:
+        b_bias = gl.load(bias + o_d, mask=o_d < D, other=0).to(gl.float32)
+    for i_r in gl.static_range(BC):
+        b_y = gl.full((BT // BC, BD), 0, gl.float32, layout)
+        for i_w in gl.static_range(W):
+            b_y += b_x[i_r + i_w] * b_w[i_w][None, :]
+        if bias is not None:
+            b_y += b_bias[None, :]
+        if ACTIVATION == 'silu' or ACTIVATION == 'swish':
+            b_y *= 1. / (1. + gl.exp(-b_y))
+        o_y = (bos + o_t[:, None] + i_r) * D + o_d[None, :]
+        m_y = (o_t[:, None] + i_r < T) & (o_d[None, :] < D)
+        if residual is not None:
+            b_y += gl.load(residual + o_y, mask=m_y, other=0).to(gl.float32)
+        gl.store(y + o_y, b_y, mask=m_y)
+@gluon.jit(do_not_specialize=['T'])
+def causal_conv1d_bwd_kernel(
+    x,
+    weight,
+    bias,
+    dy,
+    dx,
+    dw_partial,
+    db_partial,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    SXN: gl.constexpr,
+    SXT: gl.constexpr,
+    SYN: gl.constexpr,
+    SYT: gl.constexpr,
+    SYD: gl.constexpr,
+    D: gl.constexpr,
+    W: gl.constexpr,
+    BT: gl.constexpr,
+    BD: gl.constexpr,
+    BC: gl.constexpr,
+    ACTIVATION: gl.constexpr,
+):
+    layout: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, BD // 32],
+        threads_per_warp=[1, 32],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    i_d = gl.program_id(0).to(gl.int64)
+    i_t = gl.program_id(1).to(gl.int64)
+    i_b = gl.program_id(2).to(gl.int64)
+    i_tg = i_b * gl.num_programs(1) + i_t
+    if cu_seqlens is not None:
+        i_n = gl.load(chunk_indices + 2 * i_t).to(gl.int64)
+        i_t = gl.load(chunk_indices + 2 * i_t + 1).to(gl.int64)
+        bos = gl.load(cu_seqlens + i_n).to(gl.int64)
+        T = gl.load(cu_seqlens + i_n + 1).to(gl.int64) - bos
+        p_x = x + bos * SXT
+        p_dy = dy + bos * SYT
+    else:
+        bos = i_b * T
+        p_x = x + i_b * SXN
+        p_dy = dy + i_b * SYN
+    o_d = i_d * BD + gl.arange(0, BD, layout=gl.SliceLayout(0, layout)).to(gl.int64)
+    # each row unrolls BC consecutive tokens within the BT-token tile.
+    # rows (BT=64, BC=16): [0..15] [16..31] [32..47] [48..63]
+    o_t = i_t * BT + gl.arange(0, BT // BC, layout=gl.SliceLayout(1, layout)).to(gl.int64) * BC
+    b_w = ()
+    b_x = ()
+    b_dy = ()
+    b_dw = ()
+    for i_w in gl.static_range(W):
+        b_w += (gl.load(weight + o_d * W + i_w, mask=o_d < D, other=0).to(gl.float32),)
+        b_dw += (gl.full((BT // BC, BD), 0, gl.float32, layout),)
+    if bias is not None:
+        b_bias = gl.load(bias + o_d, mask=o_d < D, other=0).to(gl.float32)
+    for i_x in gl.static_range(BC + 2 * W - 2):
+        o_x = o_t + i_x - W + 1
+        b_xi = gl.load(
+            pointer=p_x + o_x[:, None] * SXT + o_d[None, :],
+            mask=(o_x[:, None] >= 0) & (o_x[:, None] < T) & (o_d[None, :] < D),
+            other=0,
+        ).to(gl.float32)
+        b_x += (b_xi,)
+    for i_r in gl.static_range(BC + W - 1):
+        b_dyi = gl.load(
+            pointer=p_dy + (o_t[:, None] + i_r) * SYT + o_d[None, :] * SYD,
+            mask=(o_t[:, None] + i_r < T) & (o_d[None, :] < D),
+            other=0,
+        ).to(gl.float32)
+        if ACTIVATION == 'silu' or ACTIVATION == 'swish':
+            b_y = gl.full((BT // BC, BD), 0, gl.float32, layout)
+            for i_w in gl.static_range(W):
+                b_y += b_x[i_r + i_w] * b_w[i_w][None, :]
+            if bias is not None:
+                b_y += b_bias[None, :]
+            # backward preserves the baseline's rounded recomputed preactivation.
+            b_y = b_y.to(x.dtype.element_ty).to(gl.float32)
+            b_ys = 1. / (1. + gl.exp(-b_y))
+            b_dyi = b_dyi * b_ys * (1 + b_y * (1 - b_ys))
+        b_dy += (b_dyi,)
+    b_db = gl.full((BT // BC, BD), 0, gl.float32, layout)
+    for i_r in gl.static_range(BC):
+        b_dx = gl.full((BT // BC, BD), 0, gl.float32, layout)
+        b_dw_new = ()
+        for i_w in gl.static_range(W):
+            b_dyi = b_dy[i_r + i_w]
+            b_dx += b_dyi * b_w[W - i_w - 1][None, :]
+            b_dw_new += (b_dw[i_w] + b_dyi * b_x[i_r + W - 1],)
+        b_dw = b_dw_new
+        b_db += b_dy[i_r]
+        gl.store(
+            pointer=dx + (bos + o_t[:, None] + i_r) * D + o_d[None, :],
+            value=b_dx,
+            mask=(o_t[:, None] + i_r < T) & (o_d[None, :] < D),
+        )
+    for i_w in gl.static_range(W):
+        gl.store(pointer=dw_partial + i_tg * D * W + o_d * W + W - i_w - 1, value=gl.sum(b_dw[i_w], axis=0), mask=o_d < D)
+    if db_partial is not None:
+        gl.store(db_partial + i_tg * D + o_d, gl.sum(b_db, axis=0), mask=o_d < D)
+@gluon.jit(do_not_specialize=['NP'])
+def causal_conv1d_bwd_kernel_dwdb(dw_partial, db_partial, dw, db, NP, D: gl.constexpr, W: gl.constexpr, BD: gl.constexpr):
+    BN: gl.constexpr = 128
+    layout: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, BD // 32],
+        threads_per_warp=[1, 32],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    i_d = gl.program_id(0).to(gl.int64)
+    o_d = i_d * BD + gl.arange(0, BD, layout=gl.SliceLayout(0, layout)).to(gl.int64)
+    b_dw = gl.full((BD,), 0, gl.float32, gl.SliceLayout(0, layout))
+    for i_n in range(0, NP, BN):
+        o_n = i_n + gl.arange(0, BN, layout=gl.SliceLayout(1, layout)).to(gl.int64)
+        b_partial = gl.load(
+            pointer=dw_partial + o_n[:, None] * D * W + o_d[None, :],
+            mask=(o_n[:, None] < NP) & (o_d[None, :] < D * W),
+            other=0,
+        )
+        b_dw += gl.sum(b_partial, axis=0)
+    gl.store(dw + o_d, b_dw.to(dw.dtype.element_ty), mask=o_d < D * W)
+    if db_partial is not None:
+        if i_d * BD < D:
+            b_db = gl.full((BD,), 0, gl.float32, gl.SliceLayout(0, layout))
+            for i_n in range(0, NP, BN):
+                o_n = i_n + gl.arange(0, BN, layout=gl.SliceLayout(1, layout)).to(gl.int64)
+                b_partial = gl.load(
+                    pointer=db_partial + o_n[:, None] * D + o_d[None, :],
+                    mask=(o_n[:, None] < NP) & (o_d[None, :] < D),
+                    other=0,
+                )
+                b_db += gl.sum(b_partial, axis=0)
+            gl.store(db + o_d, b_db.to(db.dtype.element_ty), mask=o_d < D)
+def causal_conv1d_fwd_verifier(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    activation: str | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    layout_fallback: bool = False,
+) -> tuple[bool, str | None]:
+    if torch.distributed.is_initialized():
+        return False, "Gluon convolution does not support distributed execution"
+    if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return False, "Gluon convolution supports float16, bfloat16, and float32"
+    if x.ndim != 3 or x.stride(-1) != 1:
+        return False, "Gluon convolution requires [B, T, D] with contiguous channels"
+    if weight is None or weight.shape[0] != x.shape[-1] or weight.shape[1] not in (2, 3, 4):
+        return False, "Gluon convolution requires a width of 2, 3, or 4"
+    if cu_seqlens is not None and x.shape[0] != 1:
+        return False, "Gluon packed convolution requires batch size 1"
+    if chunk_size != 64:
+        return False, "Gluon convolution requires 64-token chunk indices"
+    return True, None
+@register(causal_conv1d_fwd, backend=GluonBackend, verifier=causal_conv1d_fwd_verifier, env_var='FLA_CONV_GLUON')
+@input_guard(no_guard_contiguous=['x'])
+def causal_conv1d_fwd_gluon(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    activation: str | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    layout_fallback: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    del layout_fallback
+    BT = chunk_size
+    B, T, D = x.shape
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BT, cu_seqlens_cpu=cu_seqlens_cpu)
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    use_small_tile = B * T * D <= 1048576
+    BD, num_splits = (32, 2) if use_small_tile else (64, 1)
+    causal_conv1d_fwd_kernel[(triton.cdiv(D, BD), NT * num_splits, B)](
+        x=x,
+        y=y,
+        weight=weight,
+        bias=bias,
+        residual=residual,
+        cu_seqlens=cu_seqlens,
+        initial_state=initial_state,
+        chunk_indices=chunk_indices,
+        T=T,
+        SXN=x.stride(0),
+        SXT=x.stride(1),
+        D=D,
+        W=weight.shape[1],
+        BT=BT // num_splits,
+        BD=BD,
+        BC=8,
+        ACTIVATION=activation,
+        NUM_SPLITS=num_splits,
+        num_warps=4,
+    )
+    final_state = None
+    if output_final_state:
+        from ....modules.causal_conv1d import causal_conv1d_update_states
+        final_state = causal_conv1d_update_states(
+            x=x,
+            state_len=weight.shape[1],
+            initial_state=initial_state,
+            cu_seqlens=cu_seqlens,
+        )
+    return y, final_state
+def causal_conv1d_bwd_verifier(
+    x: torch.Tensor,
+    dy: torch.Tensor,
+    dht: torch.Tensor | None,
+    weight: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    activation: str | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    layout_fallback: bool = False,
+) -> tuple[bool, str | None]:
+    if initial_state is not None or dht is not None:
+        return False, "Gluon convolution does not support state gradients"
+    return causal_conv1d_fwd_verifier(x=x, weight=weight, cu_seqlens=cu_seqlens, chunk_size=chunk_size)
+@register(causal_conv1d_bwd, backend=GluonBackend, verifier=causal_conv1d_bwd_verifier, env_var='FLA_CONV_GLUON')
+def causal_conv1d_bwd_gluon(
+    x: torch.Tensor,
+    dy: torch.Tensor,
+    dht: torch.Tensor | None,
+    weight: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    activation: str | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    layout_fallback: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, None]:
+    del layout_fallback
+    BT = chunk_size
+    B, T, D = x.shape
+    W = weight.shape[1]
+    BD = 32
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BT, cu_seqlens_cpu=cu_seqlens_cpu)
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+    dx = torch.empty_like(x, memory_format=torch.contiguous_format)
+    dw_partial = weight.new_empty((B * NT, D, W), dtype=torch.float32)
+    db_partial = bias.new_empty((B * NT, D), dtype=torch.float32) if bias is not None else None
+    causal_conv1d_bwd_kernel[(triton.cdiv(D, BD), NT, B)](
+        x=x,
+        weight=weight,
+        bias=bias,
+        dy=dy,
+        dx=dx,
+        dw_partial=dw_partial,
+        db_partial=db_partial,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        SXN=x.stride(0),
+        SXT=x.stride(1),
+        SYN=dy.stride(0),
+        SYT=dy.stride(1),
+        SYD=dy.stride(2),
+        D=D,
+        W=W,
+        BT=BT,
+        BD=BD,
+        BC=16,
+        ACTIVATION=activation,
+        num_warps=4,
+    )
+    dw = weight.new_empty((D, W))
+    db = bias.new_empty((D,)) if bias is not None else None
+    causal_conv1d_bwd_kernel_dwdb[(triton.cdiv(D * W, BD),)](
+        dw_partial=dw_partial,
+        db_partial=db_partial,
+        dw=dw,
+        db=db,
+        NP=B * NT,
+        D=D,
+        W=W,
+        BD=BD,
+        num_warps=4,
+    )
+    dr = dy if residual is not None else None
+    return dx, dw, db, dr, None

build/torch-cuda/modules/{conv/cp/ops.py → causal_conv1d/cp.py} RENAMED Viewed

@@ -8,8 +8,8 @@
 import torch
 import torch.distributed as dist
-from ....ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd
-from ....ops.utils import prepare_chunk_indices
 class CausalConv1dFunctionCP(torch.autograd.Function):
@@ -36,28 +36,58 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
         """Prepare initial_state for CP forward pass by communicating with previous rank.
         Args:
-            x: Input tensor of shape [1, T, D]
-            weight: Weight tensor of shape [D, W]
-            cu_seqlens: Cumulative sequence lengths
-            context: CP context
-            group: Process group for communication
         Returns:
-            initial_state: Initial state tensor of shape [N, D, W] or None
         """
-        if group is None:
             return None
-        W = weight.shape[-1]  # weight: [D, W]
         D = weight.shape[0]
         initial_state = None
         if not context.is_first_rank:
-            # Non-first rank needs initial_state
             assert x.dim() == 3 and x.shape[0] == 1, f"CP requires [1, T, D], got {x.shape}"
             x_2d = x.squeeze(0)  # [T, D]
             tails = x_2d[-(W-1):].contiguous()  # [W-1, D]
             heads = conv_cp_send_recv_fwd(tails, group)  # [W-1, D]
-            # Construct initial_state: [N, D, W]
             N = len(cu_seqlens) - 1
             initial_state = torch.zeros(N, D, W, device=x.device, dtype=x.dtype)
             valid_len = min(W - 1, context.pre_num_conv_tokens)
@@ -65,10 +95,11 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
                 # heads[-valid_len:]: [valid_len, D] -> [D, valid_len]
                 initial_state[0, :, -valid_len:] = heads[-valid_len:].T
         else:
-            # First rank also needs to participate in communication (send tails)
             x_2d = x.squeeze(0)
             tails = x_2d[-(W-1):].contiguous()
-            _ = conv_cp_send_recv_fwd(tails, group)  # Send but don't use
         return initial_state
@@ -78,25 +109,46 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
         dh0: torch.Tensor | None,
         W: int,
         group: dist.ProcessGroup | None,
-        is_first_rank: bool,
-        pre_num_conv_tokens: int = 0,
     ) -> None:
         """Correct dx gradients for CP backward pass by communicating with next rank.
         Args:
-            dx: Gradient tensor to be corrected, shape [1, T, D]
-            dh0: Gradient w.r.t. initial_state, shape [N, D, W] or None
-            W: Kernel size
-            group: Process group for communication
-            is_first_rank: Whether this is the first rank in the sequence's processing chain
-            pre_num_conv_tokens: Number of tokens from the previous rank that
-                belong to the first sequence on the current rank. Must match the
-                value used in the forward pass to construct initial_state.
         """
-        if group is None:
             return
         D = dx.shape[-1]
         # dh0: [N, D, W] or None
         # We only care about the first sequence's initial_state gradient
         if dh0 is not None:
@@ -104,13 +156,13 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
             # previous rank. The forward fills only the last valid_len positions
             # of initial_state; gradients for the remaining (zero-padded) positions
             # must not flow back, otherwise they leak into unrelated sequences.
-            valid_len = min(W - 1, pre_num_conv_tokens)
             d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype)
             if valid_len > 0:
                 d_initial_state[-valid_len:] = dh0[0, :, -valid_len:].T
         else:
             # dh0 is None only when this is the first rank (no initial_state needed)
-            assert is_first_rank, "dh0 should not be None when is_first_rank=False"
             d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype)
         # Sync communication: send d_initial_state to previous rank, receive from next rank
         recv_d_init = conv_cp_send_recv_bwd(d_initial_state, group)  # [W-1, D]
@@ -128,9 +180,10 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
         cp_context: FLACPContext | None,
         chunk_size: int | None,
         backend: str = 'triton',
-    ):
-        # Import here to avoid circular dependency
-        from ....modules.conv.triton.ops import causal_conv1d_fwd
         if cp_context is None:
             raise ValueError("cp_context must be provided for CausalConv1dFunctionCP")
@@ -138,9 +191,9 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
         cu_seqlens_cpu = cp_context.cu_seqlens_cpu
         group = cp_context.group
-        # Get kernel_size
         W = weight.shape[-1]  # weight: [D, W]
-        # Prepare initial_state for CP
         initial_state = CausalConv1dFunctionCP._prepare_initial_state_for_cp(
             x=x,
             weight=weight,
@@ -149,7 +202,7 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
             group=group,
         )
-        ctx.save_for_backward(x, weight, bias, initial_state)
         ctx.activation = activation
         ctx.cu_seqlens = cu_seqlens
         ctx.cu_seqlens_cpu = cu_seqlens_cpu
@@ -157,62 +210,54 @@ class CausalConv1dFunctionCP(torch.autograd.Function):
         ctx.chunk_size = chunk_size
         ctx.group = group
         ctx.W = W
-        ctx.is_first_rank = cp_context.is_first_rank
-        ctx.pre_num_conv_tokens = cp_context.pre_num_conv_tokens
-        # Call original forward
         y, _ = causal_conv1d_fwd(
             x=x,
             weight=weight,
             bias=bias,
-            residual=None,
             initial_state=initial_state,
             output_final_state=False,
             activation=activation,
             cu_seqlens=cu_seqlens,
             cu_seqlens_cpu=cu_seqlens_cpu,
             chunk_indices=chunk_indices,
-            BT=chunk_size,
         )
         return y
     @staticmethod
-    def backward(ctx, dy: torch.Tensor):
-        # Import here to avoid circular dependency
-        from ....modules.conv.triton.ops import causal_conv1d_bwd
-        x, weight, bias, initial_state = ctx.saved_tensors
         group = ctx.group
         W = ctx.W
-        # Call original backward
-        dx, dw, db, _, dh0 = causal_conv1d_bwd(
             x=x,
             dy=dy,
             dht=None,
             weight=weight,
             bias=bias,
-            residual=None,
             initial_state=initial_state,
             activation=ctx.activation,
             cu_seqlens=ctx.cu_seqlens,
             cu_seqlens_cpu=ctx.cu_seqlens_cpu,
             chunk_indices=ctx.chunk_indices,
-            BT=ctx.chunk_size,
         )
-        # Correct dx gradients for CP
-        CausalConv1dFunctionCP._correct_dx_for_cp(
-            dx=dx,
-            dh0=dh0,
-            W=W,
-            group=group,
-            is_first_rank=ctx.is_first_rank,
-            pre_num_conv_tokens=ctx.pre_num_conv_tokens,
-        )
-        return dx, dw, db, None, None, None, None, None
 def causal_conv1d_cp(
@@ -224,35 +269,62 @@ def causal_conv1d_cp(
     cp_context: FLACPContext | None = None,
     chunk_size: int | None = None,
     backend: str = 'triton',
-):
-    """
-    Context Parallel version of causal_conv1d.
-    Automatically handles communication in CP environment:
-    - Forward: get initial_state from previous rank
-    - Backward: correct dx gradients
     Args:
-        x: Input tensor of shape [1, T, D]
-        weight: Weight tensor of shape [D, W]
-        bias: Bias tensor of shape [D] or None
-        activation: Activation function name or None
-        cu_seqlens: Cumulative sequence lengths
-        cu_seqlens_cpu: Cumulative sequence lengths on CPU
-        chunk_indices: Chunk indices for variable-length sequences
-        cp_context: CP context (required for CP mode)
     """
     if cp_context is None:
         raise ValueError("cp_context must be provided for causal_conv1d_cp")
     assert cp_context.conv1d_kernel_size is not None, "conv1d_kernel_size must be provided for causal_conv1d_cp"
     assert cp_context.cu_seqlens is not None, "cu_seqlens must be provided for causal_conv1d_cp"
-    assert backend in ['triton'], "backend must be 'triton'"
     chunk_size = chunk_size or 64
     if chunk_indices is None:
-        chunk_indices = prepare_chunk_indices(cp_context.cu_seqlens, chunk_size, cu_seqlens_cpu=cp_context.cu_seqlens_cpu)
-    return CausalConv1dFunctionCP.apply(
-        x, weight, bias, activation,
-        chunk_indices, cp_context, chunk_size, backend
-    )

 import torch
 import torch.distributed as dist
+from ...ops.cp import FLACPContext, all_gather_into_tensor, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd
+from ...ops.utils import prepare_chunk_indices
 class CausalConv1dFunctionCP(torch.autograd.Function):
         """Prepare initial_state for CP forward pass by communicating with previous rank.
         Args:
+            x (torch.Tensor):
+                Input tensor of shape [1, T, D]
+            weight (torch.Tensor):
+                Weight tensor of shape [D, W]
+            cu_seqlens (torch.Tensor | None):
+                Cumulative sequence lengths
+            context (FLACPContext):
+                CP context
+            group (dist.ProcessGroup | None):
+                Process group for communication
         Returns:
+            initial_state (torch.Tensor | None):
+                Initial state tensor of shape [N, D, W] or None
         """
+        W = weight.shape[-1]  # weight: [D, W]
+        if group is None or W == 1:
             return None
         D = weight.shape[0]
         initial_state = None
+        if context.layout == 'zigzag':
+            world_size = dist.get_world_size(group)
+            rank = dist.get_rank(group)
+            part_len = context.part_len
+            assert part_len >= W - 1, (
+                f"zigzag CP requires part_len ({part_len}) >= conv1d_kernel_size - 1 ({W - 1}); "
+                "increase total tokens or reduce world size"
+            )
+            assert x.dim() == 3 and x.shape[0] == 1, f"CP requires [1, T, D], got {x.shape}"
+            x_2d = x.squeeze(0)  # [T, D]
+            # last W-1 tokens of each chain part (front, back)
+            tails = torch.stack([x_2d[part_len - (W - 1): part_len], x_2d[-(W - 1):]])
+            gathered, _ = all_gather_into_tensor(tails, group=group)
+            # chain order: front parts of ranks 0..W-1, then back parts of ranks W-1..0
+            slots = torch.cat([gathered[:, 0], gathered[:, 1].flip(0)])
+            N = len(cu_seqlens) - 1
+            initial_state = torch.zeros(N, D, W, device=x.device, dtype=x.dtype)
+            for part, row in enumerate((0, context.front_num_seqs)):
+                if not context.is_first_by_part[part]:
+                    valid_len = min(W - 1, context.pre_num_conv_tokens_by_part[part])
+                    if valid_len > 0:
+                        chain_pos = rank if part == 0 else 2 * world_size - 1 - rank
+                        initial_state[row, :, -valid_len:] = slots[chain_pos - 1][-valid_len:].T
+            return initial_state
         if not context.is_first_rank:
+            # non-first rank needs initial_state
             assert x.dim() == 3 and x.shape[0] == 1, f"CP requires [1, T, D], got {x.shape}"
             x_2d = x.squeeze(0)  # [T, D]
             tails = x_2d[-(W-1):].contiguous()  # [W-1, D]
             heads = conv_cp_send_recv_fwd(tails, group)  # [W-1, D]
+            # construct initial_state: [N, D, W]
             N = len(cu_seqlens) - 1
             initial_state = torch.zeros(N, D, W, device=x.device, dtype=x.dtype)
             valid_len = min(W - 1, context.pre_num_conv_tokens)
                 # heads[-valid_len:]: [valid_len, D] -> [D, valid_len]
                 initial_state[0, :, -valid_len:] = heads[-valid_len:].T
         else:
+            # first rank also needs to participate in communication (send tails)
             x_2d = x.squeeze(0)
             tails = x_2d[-(W-1):].contiguous()
+            # send but don't use
+            _ = conv_cp_send_recv_fwd(tails, group)
         return initial_state
         dh0: torch.Tensor | None,
         W: int,
         group: dist.ProcessGroup | None,
+        cp_context: FLACPContext,
     ) -> None:
         """Correct dx gradients for CP backward pass by communicating with next rank.
         Args:
+            dx (torch.Tensor):
+                Gradient tensor to be corrected, shape [1, T, D]
+            dh0 (torch.Tensor | None):
+                Gradient w.r.t. initial_state, shape [N, D, W] or None
+            W (int):
+                Kernel size
+            group (dist.ProcessGroup | None):
+                Process group for communication
+            cp_context (FLACPContext):
+                CP context
         """
+        if group is None or W == 1:
             return
         D = dx.shape[-1]
+        if cp_context.layout == 'zigzag':
+            world_size = dist.get_world_size(group)
+            rank = dist.get_rank(group)
+            part_len = cp_context.part_len
+            # dh0 rows for chain-first parts stay zero: their forward initial_state was zeros
+            d_initial_state = torch.zeros(2, W - 1, D, device=dx.device, dtype=dx.dtype)
+            for part, row in enumerate((0, cp_context.front_num_seqs)):
+                if not cp_context.is_first_by_part[part]:
+                    valid_len = min(W - 1, cp_context.pre_num_conv_tokens_by_part[part])
+                    if valid_len > 0:
+                        d_initial_state[part, -valid_len:] = dh0[row, :, -valid_len:].T
+            gathered, _ = all_gather_into_tensor(d_initial_state, group=group)
+            # chain order: front parts of ranks 0..W-1, then back parts of ranks W-1..0
+            slots = torch.cat([gathered[:, 0], gathered[:, 1].flip(0)])
+            # the front part's chain successor is the next rank's front or own back part
+            dx[0, part_len - (W - 1): part_len, :].add_(slots[rank + 1])
+            # the back part's successor exists unless it is chain-last (rank 0)
+            if rank > 0:
+                dx[0, -(W - 1):, :].add_(slots[2 * world_size - rank])
+            return
         # dh0: [N, D, W] or None
         # We only care about the first sequence's initial_state gradient
         if dh0 is not None:
             # previous rank. The forward fills only the last valid_len positions
             # of initial_state; gradients for the remaining (zero-padded) positions
             # must not flow back, otherwise they leak into unrelated sequences.
+            valid_len = min(W - 1, cp_context.pre_num_conv_tokens)
             d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype)
             if valid_len > 0:
                 d_initial_state[-valid_len:] = dh0[0, :, -valid_len:].T
         else:
             # dh0 is None only when this is the first rank (no initial_state needed)
+            assert cp_context.is_first_rank, "dh0 should not be None when is_first_rank=False"
             d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype)
         # Sync communication: send d_initial_state to previous rank, receive from next rank
         recv_d_init = conv_cp_send_recv_bwd(d_initial_state, group)  # [W-1, D]
         cp_context: FLACPContext | None,
         chunk_size: int | None,
         backend: str = 'triton',
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # import here to avoid circular dependency
+        from ...modules.causal_conv1d import causal_conv1d_fwd
         if cp_context is None:
             raise ValueError("cp_context must be provided for CausalConv1dFunctionCP")
         cu_seqlens_cpu = cp_context.cu_seqlens_cpu
         group = cp_context.group
+        # get kernel_size
         W = weight.shape[-1]  # weight: [D, W]
+        # prepare initial_state for CP
         initial_state = CausalConv1dFunctionCP._prepare_initial_state_for_cp(
             x=x,
             weight=weight,
             group=group,
         )
+        ctx.save_for_backward(x, weight, bias, residual, initial_state)
         ctx.activation = activation
         ctx.cu_seqlens = cu_seqlens
         ctx.cu_seqlens_cpu = cu_seqlens_cpu
         ctx.chunk_size = chunk_size
         ctx.group = group
         ctx.W = W
+        ctx.cp_context = cp_context
+        # call original forward
         y, _ = causal_conv1d_fwd(
             x=x,
             weight=weight,
             bias=bias,
+            residual=residual,
             initial_state=initial_state,
             output_final_state=False,
             activation=activation,
             cu_seqlens=cu_seqlens,
             cu_seqlens_cpu=cu_seqlens_cpu,
             chunk_indices=chunk_indices,
+            chunk_size=chunk_size,
         )
         return y
     @staticmethod
+    def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
+        # import here to avoid circular dependency
+        from ...modules.causal_conv1d import causal_conv1d_bwd
+        x, weight, bias, residual, initial_state = ctx.saved_tensors
         group = ctx.group
         W = ctx.W
+        # call original backward
+        dx, dw, db, dr, dh0 = causal_conv1d_bwd(
             x=x,
             dy=dy,
             dht=None,
             weight=weight,
             bias=bias,
+            residual=residual,
             initial_state=initial_state,
             activation=ctx.activation,
             cu_seqlens=ctx.cu_seqlens,
             cu_seqlens_cpu=ctx.cu_seqlens_cpu,
             chunk_indices=ctx.chunk_indices,
+            chunk_size=ctx.chunk_size,
         )
+        # correct dx gradients for CP
+        CausalConv1dFunctionCP._correct_dx_for_cp(dx=dx, dh0=dh0, W=W, group=group, cp_context=ctx.cp_context)
+        return dx, dw, db, None, None, None, None, None, dr
 def causal_conv1d_cp(
     cp_context: FLACPContext | None = None,
     chunk_size: int | None = None,
     backend: str = 'triton',
+    residual: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply causal convolution with cross-rank history and gradient exchange.
     Args:
+        x (torch.Tensor):
+            Local input of shape `[1, T, D]`.
+        weight (torch.Tensor):
+            Convolution weights of shape `[D, W]`.
+        bias (torch.Tensor, Optional):
+            Bias of shape `[D]`. Default: `None`.
+        activation (str, Optional):
+            Output activation: `None`, `silu`, or `swish`. Default: `None`.
+        chunk_indices (torch.Tensor, Optional):
+            Precomputed chunk indices for the Triton backend. Default: `None`.
+        cp_context (FLACPContext, Optional):
+            Required CP context with cumulative sequence lengths and convolution width. Default: `None`.
+        chunk_size (int, Optional):
+            Triton chunk size; `None` selects 64. Default: `None`.
+        backend (str, Optional):
+            Convolution backend: `triton`, `cuda`, or `mix`. `mix` retains the Triton CP path.
+            CUDA requires contiguous layout, width 2–4, and at least `W - 1` tokens per rank. Default: `'triton'`.
+        residual (torch.Tensor, Optional):
+            Residual added to the output, of shape `[1, T, D]`. Default: `None`.
+    Returns:
+        torch.Tensor:
+            Local output of shape `[1, T, D]`.
     """
     if cp_context is None:
         raise ValueError("cp_context must be provided for causal_conv1d_cp")
     assert cp_context.conv1d_kernel_size is not None, "conv1d_kernel_size must be provided for causal_conv1d_cp"
     assert cp_context.cu_seqlens is not None, "cu_seqlens must be provided for causal_conv1d_cp"
+    if backend == 'mix':
+        backend = 'triton'
+    if backend not in ('triton', 'cuda'):
+        raise ValueError(f"Unsupported CP convolution backend: {backend!r}; expected 'triton', 'cuda', or 'mix'")
+    if backend == 'cuda':
+        from ...modules.causal_conv1d.backends.cuda.cp import causal_conv1d_cuda_cp
+        return causal_conv1d_cuda_cp(
+            x=x,
+            weight=weight,
+            bias=bias,
+            residual=residual,
+            activation=activation,
+            cp_context=cp_context,
+        )
     chunk_size = chunk_size or 64
     if chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(
+            cu_seqlens=cp_context.cu_seqlens,
+            chunk_size=chunk_size,
+            cu_seqlens_cpu=cp_context.cu_seqlens_cpu,
+        )
+    return CausalConv1dFunctionCP.apply(x, weight, bias, activation, chunk_indices, cp_context, chunk_size, backend, residual)

build/torch-cuda/modules/{conv/triton/kernels.py → causal_conv1d/ops.py} RENAMED Viewed

@@ -5,13 +5,27 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
 import torch
 import triton
 import triton.language as tl
 from einops import rearrange
-from ....ops.utils.cache import fla_cache_autotune
-from ....utils import IS_AMD, autotune_cache_kwargs, input_guard
 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32]
 STATIC_WARPS = 32 if not IS_AMD else 16
@@ -30,7 +44,7 @@ STATIC_WARPS = 32 if not IS_AMD else 16
         for BD in [16, 32, 64, 128]
         for num_warps in NUM_WARPS_AUTOTUNE
     ],
-    key=['D', 'W', 'NB'],
     **autotune_cache_kwargs,
 )
 @triton.jit
@@ -53,7 +67,6 @@ def causal_conv1d_fwd_kernel(
     BT: tl.constexpr,
     BW: tl.constexpr,
     BD: tl.constexpr,
-    NB: tl.constexpr,
     ACTIVATION: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
@@ -74,9 +87,11 @@ def causal_conv1d_fwd_kernel(
         p_x = x + tl.cast(i_b, tl.int64) * stride_x_n
     o_d = i_d * BD + tl.arange(0, BD)
     o_w = tl.arange(0, BW) + W - BW
     m_d = o_d < D
     m_w = o_w >= 0
     if HAS_WEIGHT:
         # [BD, BW]
@@ -85,33 +100,31 @@ def causal_conv1d_fwd_kernel(
     b_y = tl.zeros((BT, BD), dtype=tl.float32)
     if not USE_INITIAL_STATE:
         for i_w in tl.static_range(-W + 1, 1):
-            p_yi = tl.make_block_ptr(p_x, (T, D), (stride_x_t, stride_x_d), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
             # [BT, BD]
-            b_yi = tl.load(p_yi, boundary_check=(0, 1)).to(tl.float32)
             if HAS_WEIGHT:
                 b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1)
             b_y += b_yi
     elif i_t * BT >= W:
-        # to make Triton compiler happy, we need to copy codes
         for i_w in tl.static_range(-W + 1, 1):
-            p_yi = tl.make_block_ptr(p_x, (T, D), (stride_x_t, stride_x_d), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
             # [BT, BD]
-            b_yi = tl.load(p_yi, boundary_check=(0, 1)).to(tl.float32)
             if HAS_WEIGHT:
                 b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1)
             b_y += b_yi
     else:
-        o_t = i_t * BT + tl.arange(0, BT)
         for i_w in tl.static_range(-W + 1, 1):
             o_x = o_t + i_w
             m_x = ((o_x >= 0) & (o_x < T))[:, None] & m_d
             m_c = ((o_x + W >= 0) & (o_x < 0))[:, None] & m_d
-            b_yi = tl.load(
-                p_x + o_x[:, None] * stride_x_t + o_d * stride_x_d,
-                mask=m_x,
-                other=0
-            ).to(tl.float32)
             b_yi += tl.load(initial_state + i_n * D*W + o_d * W + (o_x + W)[:, None], mask=m_c, other=0).to(tl.float32)
@@ -126,12 +139,12 @@ def causal_conv1d_fwd_kernel(
         b_y = b_y * tl.sigmoid(b_y)
     if HAS_RESIDUAL:
-        p_residual = tl.make_block_ptr(residual + bos * D, (T, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
-        b_residual = tl.load(p_residual, boundary_check=(0, 1))
         b_y += b_residual
-    p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
-    tl.store(p_y, tl.cast(b_y, dtype=p_y.dtype.element_ty, fp_downcast_rounding='rtne'), boundary_check=(0, 1))
 @triton.heuristics({
@@ -145,9 +158,9 @@ def causal_conv1d_fwd_kernel(
     configs=[
         triton.Config({'BD': BD}, num_warps=num_warps)
         for BD in [16, 32, 64, 128]
-        for num_warps in [4, 8, 16, 32]
     ],
-    key=['D', 'W', 'NB'],
     **autotune_cache_kwargs,
 )
 @triton.jit
@@ -165,21 +178,20 @@ def causal_conv1d_bwd_kernel(
     chunk_indices,
     B,
     T,
-    stride_x_n,   # x batch stride
-    stride_x_t,   # x time stride
-    stride_x_d,   # x dim stride
-    stride_dx_n,  # dx batch stride
-    stride_dx_t,  # dx time stride
-    stride_dx_d,  # dx dim stride
-    stride_dy_n,  # dy batch stride
-    stride_dy_t,  # dy time stride
-    stride_dy_d,  # dy dim stride
     D: tl.constexpr,
     W: tl.constexpr,
     BT: tl.constexpr,
     BW: tl.constexpr,
     BD: tl.constexpr,
-    NB: tl.constexpr,
     ACTIVATION: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
@@ -206,13 +218,15 @@ def causal_conv1d_bwd_kernel(
         p_dy = dy + tl.cast(i_b, tl.int64) * stride_dy_n
     o_d = i_d * BD + tl.arange(0, BD)
     o_w = tl.arange(0, BW) + W - BW
     m_d = o_d < D
     m_w = o_w >= 0
     if HAS_WEIGHT:
-        p_x = tl.make_block_ptr(p_x, (T, D), (stride_x_t, stride_x_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
-        b_x = tl.load(p_x, boundary_check=(0, 1))
         # [BD, BW]
         b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0)
@@ -222,13 +236,13 @@ def causal_conv1d_bwd_kernel(
     if not USE_FINAL_STATE and not USE_INITIAL_STATE:
         for i_w in tl.static_range(0, W):
-            p_dy_blk = tl.make_block_ptr(p_dy, (T, D), (stride_dy_t, stride_dy_d),
-                                         (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
             # [BT, BD]
-            b_dy = tl.load(p_dy_blk, boundary_check=(0, 1)).to(tl.float32)
             if ACTIVATION == 'swish' or ACTIVATION == 'silu':
-                p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
-                b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32)
                 b_ys = tl.sigmoid(b_y)
                 b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
             b_wdy = b_dy
@@ -242,15 +256,15 @@ def causal_conv1d_bwd_kernel(
                 b_db += tl.sum(b_dy, 0)
             b_dx += b_wdy
     elif i_t * BT >= W:
-        # to make Triton compiler happy, we need to copy codes
         for i_w in tl.static_range(0, W):
-            p_dy_blk = tl.make_block_ptr(p_dy, (T, D), (stride_dy_t, stride_dy_d),
-                                         (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
             # [BT, BD]
-            b_dy = tl.load(p_dy_blk, boundary_check=(0, 1)).to(tl.float32)
             if ACTIVATION == 'swish' or ACTIVATION == 'silu':
-                p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
-                b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32)
                 b_ys = tl.sigmoid(b_y)
                 b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
             b_wdy = b_dy
@@ -264,38 +278,42 @@ def causal_conv1d_bwd_kernel(
                 b_db += tl.sum(b_dy, 0)
             b_dx += b_wdy
     else:
-        # which may use initial state
-        o_t = i_t * BT + tl.arange(0, BT)
         for i_w in tl.static_range(0, W):
-            p_dy_blk = tl.make_block_ptr(p_dy, (T, D), (stride_dy_t, stride_dy_d),
-                                         (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
-            b_dy_shift = tl.load(p_dy_blk, boundary_check=(0, 1)).to(tl.float32)
             if ACTIVATION == 'swish' or ACTIVATION == 'silu':
-                p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0))
-                b_y_shift = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32)
                 b_ys = tl.sigmoid(b_y_shift)
                 b_dy_shift = b_dy_shift * b_ys * (1 + b_y_shift * (1 - b_ys))
             if HAS_WEIGHT:
-                # gradient comes from x：sum_t dy[t+i_w] * x[t]
                 b_dw = tl.sum(b_dy_shift * b_x, 0)
-                # index of cache：c = W - i_w + t
                 if USE_INITIAL_STATE:
                     mask_head_rows = (o_t < i_w) & (o_t < T)
-                    # dy_head = dy[t]
-                    b_dy_head = tl.load(p_dy + o_t[:, None] * stride_dy_t + o_d * stride_dy_d, mask=(mask_head_rows[:, None] & m_d[None, :]),
-                                        other=0.0).to(tl.float32)
                     if ACTIVATION == 'swish' or ACTIVATION == 'silu':
-                        # use y[t] （not y[t+i_w]）
-                        b_y_head = tl.load(y + bos * D + o_t[:, None] * D + o_d,
-                                           mask=(mask_head_rows[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
                         b_ys_head = tl.sigmoid(b_y_head)
                         b_dy_head = b_dy_head * b_ys_head * (1 + b_y_head * (1 - b_ys_head))
                     o_c = W - i_w + o_t
                     # index 0 is padding 0
                     mask_c = (mask_head_rows & (o_c >= 1) & (o_c < W))
-                    b_xc = tl.load(initial_state + i_n * D * W + o_d[None, :] * W + o_c[:, None],
-                                   mask=(mask_c[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
-                    # add the gradient comes from initial_state
                     b_dw += tl.sum(b_dy_head * b_xc, 0)
                 tl.store(dw + i_tg * D * W + o_d * W + W - i_w - 1, b_dw.to(dw.dtype.element_ty), mask=m_d)
@@ -312,9 +330,11 @@ def causal_conv1d_bwd_kernel(
         if i_t * BT + BT >= T-W:
             start_tok = max(0, T - (W - 1))
             offset = i_t * BT + tl.arange(0, BT)
-            tok_idx = offset - start_tok
             mask = (offset >= start_tok) & (offset < T)
-            w_idx = 1 + tok_idx
             dht_off = i_n * D * W + o_d[None, :] * W + w_idx[:, None]
             b_dht = tl.load(dht + dht_off, mask=mask[:, None] & m_d[None, :], other=0.).to(tl.float32)
             b_dx += b_dht
@@ -324,8 +344,8 @@ def causal_conv1d_bwd_kernel(
     else:
         p_dx = dx + tl.cast(i_b, tl.int64) * stride_dx_n
-    p_dx = tl.make_block_ptr(p_dx, (T, D), (stride_dx_t, stride_dx_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
-    tl.store(p_dx, tl.cast(b_dx, dtype=p_dx.dtype.element_ty, fp_downcast_rounding='rtne'), boundary_check=(0, 1))
 @triton.heuristics({
@@ -352,10 +372,10 @@ def causal_conv1d_update_kernel(
     y,
     weight,
     bias,
-    stride_x_n,  # batch stride
-    stride_x_d,  # dim stride
-    stride_y_n,  # batch stride
-    stride_y_d,  # dim stride
     D: tl.constexpr,
     W: tl.constexpr,
     BD: tl.constexpr,
@@ -366,7 +386,9 @@ def causal_conv1d_update_kernel(
     HAS_BIAS: tl.constexpr,
     HAS_RESIDUAL: tl.constexpr,
 ):
-    i_d, i_n = tl.program_id(0), tl.program_id(1)
     o_d = i_d * BD + tl.arange(0, BD)
     o_w = tl.arange(0, BW)
@@ -379,18 +401,9 @@ def causal_conv1d_update_kernel(
     b_cache = tl.zeros((BD, BW), dtype=tl.float32)
     if USE_INITIAL_STATE:
-        # 2. Shift Cache (Read [1:])
-        p_cache_read = tl.make_block_ptr(
-            cache + i_n * D*W,
-            shape=(D, W),
-            strides=(W, 1),
-            offsets=(i_d * BD, 1),
-            block_shape=(BD, BW),
-            order=(1, 0)
-        )
-        b_cache = tl.load(p_cache_read, boundary_check=(0, 1)).to(tl.float32)
-        # 3. Fill x to the last position
         m_update = o_w == (W - 1)
         b_cache = tl.where(m_update[None, :], b_x[:, None], b_cache)
@@ -409,24 +422,24 @@ def causal_conv1d_update_kernel(
     if HAS_RESIDUAL:
         b_y += tl.load(residual + i_n * D + o_d, mask=m_d, other=0)
-    tl.store(y + i_n * stride_y_n + o_d * stride_y_d, tl.cast(b_y,
-             dtype=y.dtype.element_ty, fp_downcast_rounding='rtne'), mask=m_d)
     if USE_INITIAL_STATE:
-        p_cache_write = tl.make_block_ptr(
-            cache + i_n * D*W,
-            shape=(D, W),
-            strides=(W, 1),
-            offsets=(i_d * BD, 0),
-            block_shape=(BD, BW),
-            order=(1, 0)
         )
-        tl.store(p_cache_write, tl.cast(b_cache, dtype=cache.dtype.element_ty,
-                 fp_downcast_rounding='rtne'), boundary_check=(0, 1))
 @triton.heuristics({
     'USE_ACTIVATION': lambda args: args['y'] is not None,
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
 })
 @triton.jit
@@ -435,69 +448,75 @@ def compute_dh0_kernel(
     y,
     weight,
     dh0,
     cu_seqlens,
     stride_dy_n,
     stride_dy_t,
     T,
     D: tl.constexpr,
     W: tl.constexpr,
     BD: tl.constexpr,
     USE_ACTIVATION: tl.constexpr,
     IS_VARLEN: tl.constexpr,
 ):
     """
     Compute dh0 (gradient w.r.t. initial_state) in a separate kernel.
     This avoids Triton compiler bugs on some architectures (e.g., GB200).
-    Grid: (cdiv(D, BD), N)
     """
-    i_d, i_n = tl.program_id(0), tl.program_id(1)
-    # Get sequence boundaries
     if IS_VARLEN:
         bos = tl.load(cu_seqlens + i_n).to(tl.int64)
         eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
         seq_len = eos - bos
-        # For varlen, dy is [1, total_T, D], offset by bos
         dy_base = dy + bos * stride_dy_t
     else:
         seq_len = T
-        # For non-varlen, dy is [B, T, D], offset by i_n * stride_dy_n
         dy_base = dy + tl.cast(i_n, tl.int64) * stride_dy_n
     o_d = i_d * BD + tl.arange(0, BD)
     m_d = o_d < D
-    # For each i_w in [1, W), compute dh0[i_n, :, i_w]
     for i_w in tl.static_range(1, W):
         b_dh0 = tl.zeros([BD], dtype=tl.float32)
-        # Accumulate contributions from t = 0 to min(i_w, seq_len) - 1
         for t in tl.static_range(0, W - 1):
             if t < i_w:
                 w_idx = i_w - 1 - t
-                # Load dy[t, :] relative to dy_base
-                p_dy = dy_base + t * stride_dy_t + o_d
                 m_t = (t < seq_len) & m_d
                 b_dy = tl.load(p_dy, mask=m_t, other=0).to(tl.float32)
                 if USE_ACTIVATION:
                     if IS_VARLEN:
-                        p_y = y + bos * stride_dy_t + t * stride_dy_t + o_d
                     else:
-                        p_y = y + tl.cast(i_n, tl.int64) * stride_dy_n + t * stride_dy_t + o_d
                     b_y = tl.load(p_y, mask=m_t, other=0).to(tl.float32)
                     b_ys = tl.sigmoid(b_y)
                     b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
-                # Get weight[:, w_idx]
                 b_w_col = tl.load(weight + o_d * W + w_idx, mask=m_d, other=0).to(tl.float32)
-                # Accumulate
                 b_dh0 += tl.where(m_t, b_dy * b_w_col, 0)
-        # Store dh0[i_n, :, i_w]
         p_dh0 = dh0 + i_n * D * W + o_d * W + i_w
         tl.store(p_dh0, b_dh0.to(dh0.dtype.element_ty), mask=m_d)
@@ -523,9 +542,11 @@ def causal_conv1d_states_fwd_kernel(
     USE_INITIAL_STATE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
 ):
-    i_d, i_n = tl.program_id(0), tl.program_id(1)
-    # o_d Shape: [BD]
     o_d = i_d * BD + tl.arange(0, BD)
     m_d = o_d < D
@@ -538,10 +559,11 @@ def causal_conv1d_states_fwd_kernel(
         seq_len = T
         p_x = x + tl.cast(i_n, tl.int64) * stride_x_n
-    p_x = tl.make_block_ptr(p_x, (seq_len, D), (stride_x_t, stride_x_d), (seq_len - BW, i_d * BD), (BW, BD), (1, 0))
-    # b_x Shape: [BW, BD]
-    b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
     if USE_INITIAL_STATE:
         if seq_len < BW:
@@ -554,21 +576,245 @@ def causal_conv1d_states_fwd_kernel(
             b_cache = tl.load(p_init, mask=mask_init, other=0)
             b_x += b_cache
-    # final_state: [N, D, W] (Channel Major inside sample)
-    # o_w Shape: [BW]
     o_w = W - BW + tl.arange(0, BW)
     # o_d[:, None] -> [BD, 1]
     # o_w[None, :] -> [1, BW]
-    # p_final Shape -> [BD, BW]
     p_final = final_state + tl.cast(i_n, tl.int64) * D*W + o_d[:, None] * W + o_w[None, :]
-    # m_final Shape -> [BD, BW]
     m_final = m_d[:, None] & (o_w[None, :] >= 0)
     tl.store(p_final, tl.trans(b_x).to(final_state.dtype.element_ty), mask=m_final)
 @input_guard(no_guard_contiguous=["x"])
 def causal_conv1d_update_states(
     x: torch.Tensor,
@@ -598,7 +844,7 @@ def causal_conv1d_update_states(
     BD = min(triton.next_power_of_2(D), 256)
     BW = triton.next_power_of_2(W)
-    grid = (triton.cdiv(D, BD), N)
     causal_conv1d_states_fwd_kernel[grid](
         x=x,
@@ -611,12 +857,13 @@ def causal_conv1d_update_states(
         stride_x_n=stride_x_n,
         stride_x_t=stride_x_t,
         stride_x_d=stride_x_d,
-        BW=BW,
         BD=BD,
     )
     return final_state
 @input_guard(no_guard_contiguous=["x"])
 def causal_conv1d_update(
     x: torch.Tensor,
@@ -636,21 +883,18 @@ def causal_conv1d_update(
     BW = triton.next_power_of_2(W)
     if x.dim() == 2:
-        # Case: (N, D)
         stride_x_n = x.stride(0)
         stride_x_d = x.stride(1)
     elif x.dim() == 3 and x.shape[0] == 1:
-        # Case: (1, N, D) -> Time=1, Batch=N, Dim=D
-        # Batch 在 dim 1
         stride_x_n = x.stride(1)
         stride_x_d = x.stride(2)
     elif x.dim() == 3:
-        # Case: (N, 1, D) -> Batch=N, Time=1, Dim=D
-        # Batch 在 dim 0
         stride_x_n = x.stride(0)
         stride_x_d = x.stride(2)
     else:
-        # Fallback / Error case
         raise ValueError(f"Unsupported input shape: {x.shape}")
     y = torch.empty_like(x, memory_format=torch.contiguous_format)
@@ -662,7 +906,9 @@ def causal_conv1d_update(
     elif y.dim() == 3:
         stride_y_n, stride_y_d = y.stride(0), y.stride(2)
-    def grid(meta): return (triton.cdiv(D, meta['BD']), N)
     causal_conv1d_update_kernel[grid](
         x=x,
@@ -681,3 +927,398 @@ def causal_conv1d_update(
         ACTIVATION=activation,
     )
     return y.view(shape), cache

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+import os
+import warnings
 import torch
+import torch.nn as nn
 import triton
 import triton.language as tl
 from einops import rearrange
+from ...backends import dispatch
+from ...ops.cp import FLACPContext
+from ...ops.utils import prepare_chunk_indices
+from ...ops.utils.cache import fla_cache_autotune
+from ...utils import IS_AMD, autotune_cache_kwargs, deprecate_kwarg, input_guard
+try:
+    from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda
+    from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda
+except ImportError:
+    causal_conv1d_fn_cuda = None
+    causal_conv1d_update_cuda = None
 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32]
 STATIC_WARPS = 32 if not IS_AMD else 16
         for BD in [16, 32, 64, 128]
         for num_warps in NUM_WARPS_AUTOTUNE
     ],
+    key=['D', 'W'],
     **autotune_cache_kwargs,
 )
 @triton.jit
     BT: tl.constexpr,
     BW: tl.constexpr,
     BD: tl.constexpr,
     ACTIVATION: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
         p_x = x + tl.cast(i_b, tl.int64) * stride_x_n
     o_d = i_d * BD + tl.arange(0, BD)
+    o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT)
     o_w = tl.arange(0, BW) + W - BW
     m_d = o_d < D
     m_w = o_w >= 0
+    m_y = (o_t < T)[:, None] & m_d[None, :]
     if HAS_WEIGHT:
         # [BD, BW]
     b_y = tl.zeros((BT, BD), dtype=tl.float32)
     if not USE_INITIAL_STATE:
         for i_w in tl.static_range(-W + 1, 1):
+            o_x = o_t + i_w
+            p_yi = p_x + o_x[:, None] * stride_x_t + o_d[None, :] * stride_x_d
             # [BT, BD]
+            b_yi = tl.load(p_yi, mask=((o_x >= 0) & (o_x < T))[:, None] & m_d[None, :], other=0.0).to(tl.float32)
             if HAS_WEIGHT:
                 b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1)
             b_y += b_yi
     elif i_t * BT >= W:
+        # separate branches avoid a Triton compiler issue.
         for i_w in tl.static_range(-W + 1, 1):
+            o_x = o_t + i_w
+            p_yi = p_x + o_x[:, None] * stride_x_t + o_d[None, :] * stride_x_d
             # [BT, BD]
+            b_yi = tl.load(p_yi, mask=((o_x >= 0) & (o_x < T))[:, None] & m_d[None, :], other=0.0).to(tl.float32)
             if HAS_WEIGHT:
                 b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1)
             b_y += b_yi
     else:
+        o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT)
         for i_w in tl.static_range(-W + 1, 1):
             o_x = o_t + i_w
             m_x = ((o_x >= 0) & (o_x < T))[:, None] & m_d
             m_c = ((o_x + W >= 0) & (o_x < 0))[:, None] & m_d
+            b_yi = tl.load(p_x + o_x[:, None] * stride_x_t + o_d * stride_x_d, mask=m_x, other=0).to(tl.float32)
             b_yi += tl.load(initial_state + i_n * D*W + o_d * W + (o_x + W)[:, None], mask=m_c, other=0).to(tl.float32)
         b_y = b_y * tl.sigmoid(b_y)
     if HAS_RESIDUAL:
+        p_residual = residual + bos * D + o_t[:, None] * D + o_d[None, :]
+        b_residual = tl.load(p_residual, mask=m_y, other=0.0)
         b_y += b_residual
+    p_y = y + bos * D + o_t[:, None] * D + o_d[None, :]
+    tl.store(p_y, tl.cast(b_y, dtype=p_y.dtype.element_ty, fp_downcast_rounding='rtne'), mask=m_y)
 @triton.heuristics({
     configs=[
         triton.Config({'BD': BD}, num_warps=num_warps)
         for BD in [16, 32, 64, 128]
+        for num_warps in NUM_WARPS_AUTOTUNE
     ],
+    key=['D', 'W'],
     **autotune_cache_kwargs,
 )
 @triton.jit
     chunk_indices,
     B,
     T,
+    stride_x_n,
+    stride_x_t,
+    stride_x_d,
+    stride_dx_n,
+    stride_dx_t,
+    stride_dx_d,
+    stride_dy_n,
+    stride_dy_t,
+    stride_dy_d,
     D: tl.constexpr,
     W: tl.constexpr,
     BT: tl.constexpr,
     BW: tl.constexpr,
     BD: tl.constexpr,
     ACTIVATION: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
         p_dy = dy + tl.cast(i_b, tl.int64) * stride_dy_n
     o_d = i_d * BD + tl.arange(0, BD)
+    o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT)
     o_w = tl.arange(0, BW) + W - BW
     m_d = o_d < D
     m_w = o_w >= 0
+    m_x = (o_t < T)[:, None] & m_d[None, :]
     if HAS_WEIGHT:
+        p_x = p_x + o_t[:, None] * stride_x_t + o_d[None, :] * stride_x_d
+        b_x = tl.load(p_x, mask=m_x, other=0.0)
         # [BD, BW]
         b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0)
     if not USE_FINAL_STATE and not USE_INITIAL_STATE:
         for i_w in tl.static_range(0, W):
+            o_dy = o_t + i_w
+            p_dy_blk = p_dy + o_dy[:, None] * stride_dy_t + o_d[None, :] * stride_dy_d
             # [BT, BD]
+            b_dy = tl.load(p_dy_blk, mask=((o_dy < T)[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
             if ACTIVATION == 'swish' or ACTIVATION == 'silu':
+                p_y = y + bos * D + o_dy[:, None] * D + o_d[None, :]
+                b_y = tl.load(p_y, mask=((o_dy < T)[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
                 b_ys = tl.sigmoid(b_y)
                 b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
             b_wdy = b_dy
                 b_db += tl.sum(b_dy, 0)
             b_dx += b_wdy
     elif i_t * BT >= W:
+        # separate branches avoid a Triton compiler issue.
         for i_w in tl.static_range(0, W):
+            o_dy = o_t + i_w
+            p_dy_blk = p_dy + o_dy[:, None] * stride_dy_t + o_d[None, :] * stride_dy_d
             # [BT, BD]
+            b_dy = tl.load(p_dy_blk, mask=((o_dy < T)[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
             if ACTIVATION == 'swish' or ACTIVATION == 'silu':
+                p_y = y + bos * D + o_dy[:, None] * D + o_d[None, :]
+                b_y = tl.load(p_y, mask=((o_dy < T)[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
                 b_ys = tl.sigmoid(b_y)
                 b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
             b_wdy = b_dy
                 b_db += tl.sum(b_dy, 0)
             b_dx += b_wdy
     else:
+        o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT)
         for i_w in tl.static_range(0, W):
+            o_dy = o_t + i_w
+            p_dy_blk = p_dy + o_dy[:, None] * stride_dy_t + o_d[None, :] * stride_dy_d
+            b_dy_shift = tl.load(p_dy_blk, mask=((o_dy < T)[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
             if ACTIVATION == 'swish' or ACTIVATION == 'silu':
+                p_y = y + bos * D + o_dy[:, None] * D + o_d[None, :]
+                b_y_shift = tl.load(p_y, mask=((o_dy < T)[:, None] & m_d[None, :]), other=0.0).to(tl.float32)
                 b_ys = tl.sigmoid(b_y_shift)
                 b_dy_shift = b_dy_shift * b_ys * (1 + b_y_shift * (1 - b_ys))
             if HAS_WEIGHT:
                 b_dw = tl.sum(b_dy_shift * b_x, 0)
                 if USE_INITIAL_STATE:
                     mask_head_rows = (o_t < i_w) & (o_t < T)
+                    b_dy_head = tl.load(
+                        p_dy + o_t[:, None] * stride_dy_t + o_d * stride_dy_d,
+                        mask=mask_head_rows[:, None] & m_d[None, :],
+                        other=0.0,
+                    ).to(tl.float32)
                     if ACTIVATION == 'swish' or ACTIVATION == 'silu':
+                        # cache gradients use the unshifted preactivation y[t].
+                        b_y_head = tl.load(
+                            y + bos * D + o_t[:, None] * D + o_d,
+                            mask=mask_head_rows[:, None] & m_d[None, :],
+                            other=0.0,
+                        ).to(tl.float32)
                         b_ys_head = tl.sigmoid(b_y_head)
                         b_dy_head = b_dy_head * b_ys_head * (1 + b_y_head * (1 - b_ys_head))
                     o_c = W - i_w + o_t
                     # index 0 is padding 0
                     mask_c = (mask_head_rows & (o_c >= 1) & (o_c < W))
+                    b_xc = tl.load(
+                        initial_state + i_n * D * W + o_d[None, :] * W + o_c[:, None],
+                        mask=mask_c[:, None] & m_d[None, :],
+                        other=0.0,
+                    ).to(tl.float32)
                     b_dw += tl.sum(b_dy_head * b_xc, 0)
                 tl.store(dw + i_tg * D * W + o_d * W + W - i_w - 1, b_dw.to(dw.dtype.element_ty), mask=m_d)
         if i_t * BT + BT >= T-W:
             start_tok = max(0, T - (W - 1))
             offset = i_t * BT + tl.arange(0, BT)
             mask = (offset >= start_tok) & (offset < T)
+            # state slot w is token T - W + w, so token t takes dht[:, t + W - T].
+            # start_tok is clamped, which only matters for the mask (offset >= 0 anyway),
+            # so the slot index has to come from the unclamped expression.
+            w_idx = offset + (W - T)
             dht_off = i_n * D * W + o_d[None, :] * W + w_idx[:, None]
             b_dht = tl.load(dht + dht_off, mask=mask[:, None] & m_d[None, :], other=0.).to(tl.float32)
             b_dx += b_dht
     else:
         p_dx = dx + tl.cast(i_b, tl.int64) * stride_dx_n
+    p_dx = p_dx + o_t[:, None] * stride_dx_t + o_d[None, :] * stride_dx_d
+    tl.store(p_dx, tl.cast(b_dx, dtype=p_dx.dtype.element_ty, fp_downcast_rounding='rtne'), mask=m_x)
 @triton.heuristics({
     y,
     weight,
     bias,
+    stride_x_n,
+    stride_x_d,
+    stride_y_n,
+    stride_y_d,
     D: tl.constexpr,
     W: tl.constexpr,
     BD: tl.constexpr,
     HAS_BIAS: tl.constexpr,
     HAS_RESIDUAL: tl.constexpr,
 ):
+    pid = tl.program_id(0)
+    ND = tl.cdiv(D, BD)
+    i_d, i_n = pid % ND, (pid // ND).to(tl.int64)
     o_d = i_d * BD + tl.arange(0, BD)
     o_w = tl.arange(0, BW)
     b_cache = tl.zeros((BD, BW), dtype=tl.float32)
     if USE_INITIAL_STATE:
+        p_cache_read = cache + i_n * D*W + o_d[:, None] * W + (o_w + 1)[None, :]
+        b_cache = tl.load(p_cache_read, mask=m_d[:, None] & ((o_w + 1) < W)[None, :], other=0.0).to(tl.float32)
         m_update = o_w == (W - 1)
         b_cache = tl.where(m_update[None, :], b_x[:, None], b_cache)
     if HAS_RESIDUAL:
         b_y += tl.load(residual + i_n * D + o_d, mask=m_d, other=0)
+    tl.store(
+        y + i_n * stride_y_n + o_d * stride_y_d,
+        tl.cast(b_y, dtype=y.dtype.element_ty, fp_downcast_rounding='rtne'),
+        mask=m_d,
+    )
     if USE_INITIAL_STATE:
+        p_cache_write = cache + i_n * D*W + o_d[:, None] * W + o_w[None, :]
+        tl.store(
+            p_cache_write,
+            tl.cast(b_cache, dtype=cache.dtype.element_ty, fp_downcast_rounding='rtne'),
+            mask=m_d[:, None] & m_w[None, :],
         )
 @triton.heuristics({
     'USE_ACTIVATION': lambda args: args['y'] is not None,
+    'USE_FINAL_STATE': lambda args: args['dht'] is not None,
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
 })
 @triton.jit
     y,
     weight,
     dh0,
+    dht,
     cu_seqlens,
     stride_dy_n,
     stride_dy_t,
+    stride_dy_d,
     T,
     D: tl.constexpr,
     W: tl.constexpr,
     BD: tl.constexpr,
     USE_ACTIVATION: tl.constexpr,
+    USE_FINAL_STATE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
 ):
     """
     Compute dh0 (gradient w.r.t. initial_state) in a separate kernel.
     This avoids Triton compiler bugs on some architectures (e.g., GB200).
+    Grid: (cdiv(D, BD) * N,)
     """
+    pid = tl.program_id(0)
+    ND = tl.cdiv(D, BD)
+    i_d, i_n = pid % ND, (pid // ND).to(tl.int64)
     if IS_VARLEN:
         bos = tl.load(cu_seqlens + i_n).to(tl.int64)
         eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
         seq_len = eos - bos
+        # for varlen, dy is [1, total_T, D], offset by bos
         dy_base = dy + bos * stride_dy_t
     else:
         seq_len = T
+        # for non-varlen, dy is [B, T, D], offset by i_n * stride_dy_n
         dy_base = dy + tl.cast(i_n, tl.int64) * stride_dy_n
     o_d = i_d * BD + tl.arange(0, BD)
     m_d = o_d < D
     for i_w in tl.static_range(1, W):
         b_dh0 = tl.zeros([BD], dtype=tl.float32)
+        if USE_FINAL_STATE:
+            # short sequences retain initial_state[:, i_w] in final_state[:, i_w - seq_len],
+            # so that slot's gradient passes straight through.
+            if i_w >= seq_len:
+                p_dht = dht + i_n * D * W + o_d * W + (i_w - seq_len)
+                b_dh0 += tl.load(p_dht, mask=m_d, other=0).to(tl.float32)
         for t in tl.static_range(0, W - 1):
             if t < i_w:
                 w_idx = i_w - 1 - t
+                p_dy = dy_base + t * stride_dy_t + o_d * stride_dy_d
                 m_t = (t < seq_len) & m_d
                 b_dy = tl.load(p_dy, mask=m_t, other=0).to(tl.float32)
                 if USE_ACTIVATION:
+                    # `y` is recomputed contiguous [*, T, D]; only `dy` may carry other strides
                     if IS_VARLEN:
+                        p_y = y + bos * D + t * D + o_d
                     else:
+                        p_y = y + tl.cast(i_n, tl.int64) * T * D + t * D + o_d
                     b_y = tl.load(p_y, mask=m_t, other=0).to(tl.float32)
                     b_ys = tl.sigmoid(b_y)
                     b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys))
                 b_w_col = tl.load(weight + o_d * W + w_idx, mask=m_d, other=0).to(tl.float32)
                 b_dh0 += tl.where(m_t, b_dy * b_w_col, 0)
         p_dh0 = dh0 + i_n * D * W + o_d * W + i_w
         tl.store(p_dh0, b_dh0.to(dh0.dtype.element_ty), mask=m_d)
     USE_INITIAL_STATE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
 ):
+    pid = tl.program_id(0)
+    ND = tl.cdiv(D, BD)
+    i_d, i_n = pid % ND, (pid // ND).to(tl.int64)
+    # [BD]
     o_d = i_d * BD + tl.arange(0, BD)
     m_d = o_d < D
         seq_len = T
         p_x = x + tl.cast(i_n, tl.int64) * stride_x_n
+    o_x = tl.cast(seq_len, tl.int64) - BW + tl.arange(0, BW)
+    p_x = p_x + o_x[:, None] * stride_x_t + o_d[None, :] * stride_x_d
+    # [BW, BD]
+    b_x = tl.load(p_x, mask=((o_x >= 0) & (o_x < seq_len))[:, None] & m_d[None, :], other=0.0).to(tl.float32)
     if USE_INITIAL_STATE:
         if seq_len < BW:
             b_cache = tl.load(p_init, mask=mask_init, other=0)
             b_x += b_cache
+    # final_state: [N, D, W]
+    # [BW]
     o_w = W - BW + tl.arange(0, BW)
     # o_d[:, None] -> [BD, 1]
     # o_w[None, :] -> [1, BW]
+    # [BD, BW]
     p_final = final_state + tl.cast(i_n, tl.int64) * D*W + o_d[:, None] * W + o_w[None, :]
+    # [BD, BW]
     m_final = m_d[:, None] & (o_w[None, :] >= 0)
     tl.store(p_final, tl.trans(b_x).to(final_state.dtype.element_ty), mask=m_final)
+def _has_non_standard_layout(x: torch.Tensor) -> bool:
+    """QKV-style views (stride_t != D) break triton-ascend masked kernels."""
+    if x.dtype not in (torch.float16, torch.bfloat16):
+        return False
+    if x.dim() != 3:
+        return not x.is_contiguous()
+    _, stride_t, stride_d = x.stride()
+    return stride_d == 1 and stride_t != x.shape[-1]
+@deprecate_kwarg('BT', version='0.7.0', new_name='chunk_size')
+@dispatch
+@input_guard(no_guard_contiguous=["x"])
+def causal_conv1d_fwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    activation: str | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    layout_fallback: bool = False,
+) -> torch.Tensor:
+    BT = chunk_size
+    shape = x.shape
+    if x.shape[-1] != weight.shape[0]:
+        x = rearrange(x, 'b t ... -> b t (...)')
+    B, T, D = x.shape[0], x.shape[1], weight.shape[0]
+    W = weight.shape[1]
+    stride_x_n, stride_x_t, stride_x_d = x.stride()
+    BW = triton.next_power_of_2(W)
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    def grid(meta) -> tuple[int, ...]:
+        return (triton.cdiv(D, meta['BD']), NT, B)
+    causal_conv1d_fwd_kernel[grid](
+        x=x,
+        y=y,
+        weight=weight,
+        bias=bias,
+        residual=residual,
+        cu_seqlens=cu_seqlens,
+        initial_state=initial_state,
+        chunk_indices=chunk_indices,
+        B=B,
+        T=T,
+        stride_x_n=stride_x_n,
+        stride_x_t=stride_x_t,
+        stride_x_d=stride_x_d,
+        D=D,
+        W=W,
+        BT=BT,
+        BW=BW,
+        ACTIVATION=activation,
+    )
+    final_state = None
+    if output_final_state:
+        final_state = causal_conv1d_update_states(x=x, state_len=W, initial_state=initial_state, cu_seqlens=cu_seqlens)
+    return y.view(shape), final_state
+@dispatch
+def compute_dh0_triton(
+    dy: torch.Tensor,
+    y: torch.Tensor | None,
+    weight: torch.Tensor,
+    initial_state: torch.Tensor,
+    activation: str | None,
+    cu_seqlens: torch.Tensor | None,
+    dht: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    Compute dh0 (gradient w.r.t. initial_state) using a separate Triton kernel.
+    This is a workaround for Triton compiler bugs on some architectures (e.g., GB200).
+    """
+    D, W = weight.shape
+    N = initial_state.shape[0]
+    T = dy.shape[1]
+    dh0 = torch.zeros_like(initial_state)
+    BD = 32
+    grid = (triton.cdiv(D, BD) * N,)
+    y_to_pass = y if activation in ('swish', 'silu') else None
+    # dy is [B, T, D] but may be a strided view; `y` is always contiguous
+    stride_dy_n, stride_dy_t, stride_dy_d = dy.stride()
+    compute_dh0_kernel[grid](
+        dy=dy,
+        y=y_to_pass,
+        weight=weight,
+        dh0=dh0,
+        dht=dht,
+        cu_seqlens=cu_seqlens,
+        stride_dy_n=stride_dy_n,
+        stride_dy_t=stride_dy_t,
+        stride_dy_d=stride_dy_d,
+        T=T,
+        D=D,
+        W=W,
+        BD=BD,
+    )
+    return dh0
+@deprecate_kwarg('BT', version='0.7.0', new_name='chunk_size')
+@dispatch
+def causal_conv1d_bwd(
+    x: torch.Tensor,
+    dy: torch.Tensor,
+    dht: torch.Tensor,
+    weight: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    activation: str | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    layout_fallback: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
+    BT = chunk_size
+    shape = x.shape
+    if x.shape[-1] != weight.shape[0]:
+        x = rearrange(x, 'b t ... -> b t (...)')
+    B, T, D = x.shape
+    W = weight.shape[1] if weight is not None else None
+    stride_x_n, stride_x_t, stride_x_d = x.stride()
+    stride_dy_n, stride_dy_t, stride_dy_d = dy.stride()
+    BW = triton.next_power_of_2(W)
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+    y = None
+    if activation is not None:
+        y, _ = causal_conv1d_fwd(
+            x=x,
+            weight=weight,
+            bias=bias,
+            residual=None,
+            initial_state=initial_state,
+            output_final_state=False,
+            activation=None,
+            cu_seqlens=cu_seqlens,
+            cu_seqlens_cpu=cu_seqlens_cpu,
+            chunk_indices=chunk_indices,
+        )
+    dx = torch.empty_like(x)
+    dw = weight.new_empty(B*NT, *weight.shape, dtype=torch.float) if weight is not None else None
+    db = bias.new_empty(B*NT, *bias.shape, dtype=torch.float) if bias is not None else None
+    dr = dy if residual is not None else None
+    stride_dx_n, stride_dx_t, stride_dx_d = dx.stride()
+    def grid(meta) -> tuple[int, ...]:
+        return (triton.cdiv(D, meta['BD']), NT, B)
+    causal_conv1d_bwd_kernel[grid](
+        x=x,
+        y=y,
+        weight=weight,
+        initial_state=initial_state,
+        dht=dht,
+        dy=dy,
+        dx=dx,
+        dw=dw,
+        db=db,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        B=B,
+        T=T,
+        stride_x_n=stride_x_n,
+        stride_x_t=stride_x_t,
+        stride_x_d=stride_x_d,
+        stride_dx_n=stride_dx_n,
+        stride_dx_t=stride_dx_t,
+        stride_dx_d=stride_dx_d,
+        stride_dy_n=stride_dy_n,
+        stride_dy_t=stride_dy_t,
+        stride_dy_d=stride_dy_d,
+        D=D,
+        W=W,
+        BT=BT,
+        BW=BW,
+        ACTIVATION=activation,
+    )
+    if weight is not None:
+        dw = dw.sum(0).to(weight)
+    if bias is not None:
+        db = db.sum(0).to(bias)
+    # compute dh0 using separate Triton kernel to avoid compiler bugs on some architectures (e.g., GB200)
+    dh0 = None
+    if initial_state is not None:
+        dh0 = compute_dh0_triton(
+            dy=dy,
+            y=y,
+            weight=weight,
+            initial_state=initial_state,
+            activation=activation,
+            cu_seqlens=cu_seqlens,
+            dht=dht,
+        )
+    return dx.view(shape), dw, db, dr, dh0
+@dispatch
 @input_guard(no_guard_contiguous=["x"])
 def causal_conv1d_update_states(
     x: torch.Tensor,
     BD = min(triton.next_power_of_2(D), 256)
     BW = triton.next_power_of_2(W)
+    grid = (triton.cdiv(D, BD) * N,)
     causal_conv1d_states_fwd_kernel[grid](
         x=x,
         stride_x_n=stride_x_n,
         stride_x_t=stride_x_t,
         stride_x_d=stride_x_d,
         BD=BD,
+        BW=BW,
     )
     return final_state
+@dispatch
 @input_guard(no_guard_contiguous=["x"])
 def causal_conv1d_update(
     x: torch.Tensor,
     BW = triton.next_power_of_2(W)
     if x.dim() == 2:
+        # case: (N, D)
         stride_x_n = x.stride(0)
         stride_x_d = x.stride(1)
     elif x.dim() == 3 and x.shape[0] == 1:
+        # [1, N, D]: one token per sequence
         stride_x_n = x.stride(1)
         stride_x_d = x.stride(2)
     elif x.dim() == 3:
+        # [N, 1, D]: one token per sequence
         stride_x_n = x.stride(0)
         stride_x_d = x.stride(2)
     else:
         raise ValueError(f"Unsupported input shape: {x.shape}")
     y = torch.empty_like(x, memory_format=torch.contiguous_format)
     elif y.dim() == 3:
         stride_y_n, stride_y_d = y.stride(0), y.stride(2)
+    def grid(meta) -> tuple[int, ...]:
+        return (triton.cdiv(D, meta['BD']) * N,)
     causal_conv1d_update_kernel[grid](
         x=x,
         ACTIVATION=activation,
     )
     return y.view(shape), cache
+class CausalConv1dFunction(torch.autograd.Function):
+    @staticmethod
+    @input_guard(no_guard_contiguous=["x"])
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        weight: torch.Tensor | None = None,
+        bias: torch.Tensor | None = None,
+        residual: torch.Tensor | None = None,
+        initial_state: torch.Tensor | None = None,
+        output_final_state: bool | None = False,
+        activation: str | None = None,
+        cu_seqlens: torch.Tensor | None = None,
+        cu_seqlens_cpu: torch.LongTensor | None = None,
+        chunk_indices: torch.LongTensor | None = None,
+        chunk_size: int = 64,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if cu_seqlens is not None and chunk_indices is None:
+            chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu)
+        ctx.activation = activation
+        ctx.cu_seqlens = cu_seqlens
+        ctx.cu_seqlens_cpu = cu_seqlens_cpu
+        ctx.chunk_indices = chunk_indices
+        ctx.layout_fallback = _has_non_standard_layout(x=x)
+        ctx.save_for_backward(x, weight, bias, residual, initial_state)
+        y, final_state = causal_conv1d_fwd(
+            x=x,
+            weight=weight,
+            bias=bias,
+            residual=residual,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            activation=activation,
+            cu_seqlens=cu_seqlens,
+            cu_seqlens_cpu=cu_seqlens_cpu,
+            chunk_indices=chunk_indices,
+            chunk_size=chunk_size,
+            layout_fallback=ctx.layout_fallback,
+        )
+        return y, final_state
+    @staticmethod
+    @input_guard(no_guard_contiguous=["dy"])
+    def backward(ctx, dy: torch.Tensor, dht: torch.Tensor | None = None) -> tuple[torch.Tensor | None, ...]:
+        x, weight, bias, residual, initial_state = ctx.saved_tensors
+        dx, dw, db, dr, dh0 = causal_conv1d_bwd(
+            x=x,
+            dy=dy,
+            dht=dht,
+            weight=weight,
+            bias=bias,
+            residual=residual,
+            initial_state=initial_state,
+            activation=ctx.activation,
+            cu_seqlens=ctx.cu_seqlens,
+            cu_seqlens_cpu=ctx.cu_seqlens_cpu,
+            chunk_indices=ctx.chunk_indices,
+            layout_fallback=ctx.layout_fallback,
+        )
+        return dx, dw, db, dr, dh0, None, None, None, None, None, None
+@input_guard(no_guard_contiguous=["x"])
+def causal_conv1d(
+    x: torch.Tensor,
+    weight: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool | None = False,
+    activation: str | None = None,
+    backend: str | None = 'triton',
+    cu_seqlens: torch.Tensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    cp_context: FLACPContext | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """
+    A causal 1D convolution implementation that powers Mamba/Mamba2 and DeltaNet architectures.
+    With a residual connection, this implements the Canon operation:
+    https://papers.ssrn.com/sol3/papers.cfm?abstract_id=5240330.
+    Args:
+        x (torch.Tensor):
+            Input tensor of shape [B, T, D].
+        weight (torch.Tensor, Optional):
+            Weight tensor of shape [D, W]. Default: `None`.
+        bias (torch.Tensor, Optional):
+            Bias tensor of shape [D]. Default: `None`.
+        residual (torch.Tensor, Optional):
+            Residual tensor of shape [B, T, D]. Default: `None`.
+        initial_state (torch.Tensor, Optional):
+            Initial state tensor of shape [N, D, W],
+            where `N` is the number of sequences in the batch and `W` is the kernel size.
+            If provided, the initial state is used to initialize the cache. Default: `None`.
+        output_final_state (bool, Optional):
+            Whether to output the final state of shape [N, D, W]. Default: `False`.
+        activation (str, Optional):
+            Activations applied to output, only `swish`/`silu` or `None` (i.e., no activation) are supported.
+            Default: `None`.
+        backend (str, Optional):
+            Specifies the backend to use for the convolution operation. Supported values are `'cuda'`, `'triton'`, and `'mix'`.
+            Default: `'triton'`.
+        cu_seqlens (torch.Tensor, Optional):
+            Cumulative sequence lengths. Default: `None`.
+        cu_seqlens_cpu (torch.LongTensor, Optional):
+            CPU copy of the cumulative sequence lengths. Default: `None`.
+        chunk_indices (torch.Tensor, Optional):
+            Chunk indices for variable-length sequences. Default: `None`.
+        cp_context (FLACPContext, Optional):
+            Context-parallel sequence metadata. CP does not support initial or final states.
+            CUDA CP requires contiguous layout, width 2–4, and at least `W - 1` tokens per rank.
+            The `mix` backend uses Triton in CP mode. Default: `None`.
+    Returns:
+        tuple[torch.Tensor, torch.Tensor | None]:
+            Output and final state. If `output_final_state` is `False`, the final state is `None`.
+    """
+    # import here to avoid circular dependencies
+    from ...modules.causal_conv1d.backends.cuda import causal_conv1d_cuda, fast_causal_conv1d_fn
+    from ...modules.causal_conv1d.cp import causal_conv1d_cp
+    if cp_context is not None:
+        assert initial_state is None, "Initial state is not supported for CP"
+        assert output_final_state is False, "Output final state is not supported for CP"
+        output = causal_conv1d_cp(
+            x=x,
+            weight=weight,
+            bias=bias,
+            activation=activation,
+            chunk_indices=chunk_indices,
+            cp_context=cp_context,
+            backend=backend,
+            residual=residual,
+        )
+        return output, None
+    if backend == 'triton':
+        y, final_state = CausalConv1dFunction.apply(
+            x,
+            weight,
+            bias,
+            residual,
+            initial_state,
+            output_final_state,
+            activation,
+            cu_seqlens,
+            cu_seqlens_cpu,
+            chunk_indices,
+        )
+        return y, final_state
+    elif backend == 'mix':
+        seq_idx = kwargs.get('seq_idx')
+        return fast_causal_conv1d_fn(
+            x=x,
+            weight=weight,
+            bias=bias,
+            residual=residual,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            activation=activation,
+            cu_seqlens=cu_seqlens,
+            cu_seqlens_cpu=cu_seqlens_cpu,
+            chunk_indices=chunk_indices,
+            seq_idx=seq_idx,
+        )
+    elif backend == 'cuda':
+        return causal_conv1d_cuda(
+            x=x,
+            weight=weight,
+            bias=bias,
+            residual=residual,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            activation=activation,
+            cu_seqlens=cu_seqlens,
+            cu_seqlens_cpu=cu_seqlens_cpu,
+            **kwargs,
+        )
+    else:
+        raise ValueError(f"Unsupported backend: {backend}")
+class ShortConvolution(nn.Conv1d):
+    """Depthwise causal convolution with optional activation and state caching.
+    Args:
+        hidden_size (int):
+            Number of input and output channels.
+        kernel_size (int):
+            Width of the convolution kernel.
+        bias (bool, Optional):
+            Whether to use a channel bias. Default: `False`.
+        activation (str, Optional):
+            Activation applied to the convolution output. Default: `"silu"`.
+        backend (str, Optional):
+            Backend selected unless overridden by `FLA_CONV_BACKEND`. Default: `"triton"`.
+        device (torch.device, Optional):
+            Device for the parameters. Default: `None`.
+        dtype (torch.dtype, Optional):
+            Parameter dtype. Default: `None`.
+        **kwargs:
+            Accepts the deprecated `use_fast_conv1d` option.
+    """
+    def __init__(
+        self,
+        hidden_size: int,
+        kernel_size: int,
+        bias: bool = False,
+        activation: str | None = 'silu',
+        backend: str | None = 'triton',
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            in_channels=hidden_size,
+            out_channels=hidden_size,
+            kernel_size=kernel_size,
+            padding=kernel_size - 1,
+            groups=hidden_size,
+            bias=bias,
+            device=device,
+            dtype=dtype,
+        )
+        self.hidden_size = hidden_size
+        self.activation = None
+        if activation is not None:
+            assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet."
+            self.activation = activation
+        if 'use_fast_conv1d' in kwargs:
+            warnings.warn(
+                "The `use_fast_conv1d` parameter is deprecated and will be ignored. "
+                "Please use the `backend` parameter instead.",
+            )
+        self.backend = os.environ.get('FLA_CONV_BACKEND', backend)
+        if self.backend not in ['cuda', 'triton']:
+            raise ValueError(f"Invalid backend: {self.backend}, must be one of ['cuda', 'triton']")
+        if self.backend == 'cuda':
+            if causal_conv1d_fn_cuda is None:
+                warnings.warn(
+                    "The `backend` parameter is set to `cuda`, but `causal_conv1d_fn` is not available. "
+                    "Switching to the Triton implementation instead. "
+                    "Consider installing `causal_conv1d` to enable the CUDA backend.",
+                )
+                self.backend = 'triton'
+    def extra_repr(self) -> str:
+        s = '{in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}'
+        if self.padding != (0,) * len(self.padding):
+            s += ', padding={padding}'
+        if self.dilation != (1,) * len(self.dilation):
+            s += ', dilation={dilation}'
+        if self.output_padding != (0,) * len(self.output_padding):
+            s += ', output_padding={output_padding}'
+        if self.groups != 1:
+            s += ', groups={groups}'
+        if self.bias is None:
+            s += ', bias=False'
+        if self.padding_mode != 'zeros':
+            s += ', padding_mode={padding_mode}'
+        if self.activation is not None:
+            s += ', activation={activation}'
+        s += f', backend={self.backend}'
+        return s.format(**self.__dict__)
+    def forward(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        cache: torch.Tensor | None = None,
+        output_final_state: bool = False,
+        cu_seqlens: torch.LongTensor | None = None,
+        chunk_indices: torch.LongTensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Args:
+            x (torch.Tensor):
+                Tensor of shape `[B, T, D]`. `B` must be 1 if `cu_seqlens` is provided.
+            residual (torch.Tensor, Optional):
+                Residual tensor of shape `[B, T, D]`. Default: `None`.
+            mask (torch.Tensor, Optional):
+                Attention mask for padded positions. Default: `None`.
+            cache (torch.Tensor, Optional):
+                Previous cache of shape `[N, D, W]`, where `W` is the kernel size. Default: `None`.
+                Single-token calls without autograd update it **inplace**.
+                Other calls return the updated state when `output_final_state=True`.
+            output_final_state (bool, Optional):
+                Whether to output the final state of shape `[N, D, W]`. Default: `False`.
+            cu_seqlens (torch.Tensor, Optional):
+                Cumulative sequence lengths of shape `[N+1]`. Default: `None`.
+            chunk_indices (torch.Tensor, Optional):
+                Chunk indices for variable-length sequences. Default: `None`.
+        Returns:
+            y (torch.Tensor):
+                Convolution output of shape `[B, T, D]`.
+            final_state (torch.Tensor, Optional):
+                Cache of shape `[N, D, W]` when requested, otherwise `None`.
+        """
+        B, T, *_ = x.shape
+        N = B if cu_seqlens is None else len(cu_seqlens) - 1
+        if mask is not None:
+            if cu_seqlens is not None:
+                raise ValueError("`mask` and `cu_seqlens` cannot be provided at the same time")
+            x = x.mul_(mask.unsqueeze(-1))
+        # the in-place decode kernel requires autograd off and exactly one token per sequence, including packed inputs.
+        if not torch.is_grad_enabled() and B * T == N and (cu_seqlens is None or bool((cu_seqlens.diff() == 1).all())):
+            y, cache = self.step(
+                x=x,
+                residual=residual,
+                cache=cache,
+                output_final_state=output_final_state,
+                cu_seqlens=cu_seqlens,
+            )
+            return y, cache
+        # use Triton for stateful calls unsupported by the CUDA backend.
+        if self.backend == 'cuda' and cache is not None:
+            warnings.warn(
+                "The CUDA backend does not support both `cu_seqlens` and `cache` being provided, "
+                "or both `cu_seqlens` and `output_final_state` being provided. "
+                "Switching to the Triton backend instead. ",
+                stacklevel=2,
+            )
+            self.backend = 'triton'
+        return causal_conv1d(
+            x=x,
+            weight=rearrange(self.weight, "d 1 w -> d w"),
+            bias=self.bias,
+            residual=residual,
+            initial_state=cache,
+            output_final_state=output_final_state,
+            activation=self.activation,
+            backend=self.backend,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            **kwargs,
+        )
+    def step(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None,
+        cache: torch.Tensor | None,
+        output_final_state: bool = False,
+        cu_seqlens: torch.LongTensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        B, _, D, W = *x.shape, self.kernel_size[0]
+        N = B if cu_seqlens is None else len(cu_seqlens) - 1
+        # the update kernel requires a cache even when the final state is not requested.
+        if cache is None:
+            cache = x.new_zeros(N, D, W)
+        if self.backend == 'triton':
+            y, cache = causal_conv1d_update(
+                x=x,
+                cache=cache,
+                residual=residual,
+                weight=rearrange(self.weight, "d 1 w -> d w"),
+                bias=self.bias,
+                activation=self.activation,
+            )
+            return y, (cache if output_final_state else None)
+        shape = x.shape
+        x = x.squeeze(0) if cu_seqlens is not None else x.squeeze(1)
+        # the CUDA update mutates cache in place.
+        y = causal_conv1d_update_cuda(
+            x=x,
+            conv_state=cache,
+            weight=rearrange(self.weight, "d 1 w -> d w"),
+            bias=self.bias,
+            activation=self.activation,
+        )
+        y = y.view(shape)
+        if residual is not None:
+            y.add_(residual)
+        return y, (cache if output_final_state else None)
+    @property
+    def state_size(self) -> int:
+        return self.hidden_size * self.kernel_size[0]

build/torch-cuda/modules/conv/causal_conv1d.py DELETED Viewed

@@ -1,129 +0,0 @@
-# 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
-"""Main interface for causal 1D convolution operations."""
-import torch
-from ...ops.cp import FLACPContext
-from ...utils import input_guard
-@input_guard(no_guard_contiguous=["x"])
-def causal_conv1d(
-    x: torch.Tensor,
-    weight: torch.Tensor | None = None,
-    bias: torch.Tensor | None = None,
-    residual: torch.Tensor | None = None,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool | None = False,
-    activation: str | None = None,
-    backend: str | None = 'triton',
-    cu_seqlens: torch.Tensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    chunk_indices: torch.LongTensor | None = None,
-    cp_context: FLACPContext | None = None,
-    **kwargs,
-):
-    """
-    A causal 1D convolution implementation that powers Mamba/Mamba2 and DeltaNet architectures.
-    When a residual connection is provided, this implements the Canon operation
-    described in the paper at https://papers.ssrn.com/sol3/papers.cfm?abstract_id=5240330.
-    Args:
-        x (torch.Tensor):
-            Input tensor of shape [B, T, D].
-        weight (Optional[torch.Tensor]):
-            Weight tensor of shape [D, W]. Default: `None`.
-        bias (Optional[torch.Tensor]):
-            Bias tensor of shape [D]. Default: `None`.
-        residual (Optional[torch.Tensor]):
-            Residual tensor of shape [B, T, D]. Default: `None`.
-        initial_state (Optional[torch.Tensor]):
-            Initial state tensor of shape [N, D, W],
-            where `N` is the number of sequences in the batch and `W` is the kernel size.
-            If provided, the initial state is used to initialize the cache. Default: `None`.
-        output_final_state (Optional[bool]):
-            Whether to output the final state of shape [N, D, W]. Default: `False`.
-        activation (Optional[str]):
-            Activations applied to output, only `swish`/`silu` or `None` (i.e., no activation) are supported.
-            Default: `None`.
-        backend (Optional[str]):
-            Specifies the backend to use for the convolution operation. Supported values are `'cuda'` 、 `'triton'` and `'mix'`.
-            Default: `'triton'`.
-        cu_seqlens (Optional[torch.Tensor]):
-            Cumulative sequence lengths (optional)
-        chunk_indices (Optional[torch.LongTensor]):
-            Chunk indices for variable-length sequences (optional)
-    Returns:
-        Tuple of (output, final_state).
-        If `output_final_state` is `False`, the final state is `None`.
-    """
-    # Import here to avoid circular dependencies
-    from ...modules.conv.cp import causal_conv1d_cp
-    from ...modules.conv.cuda import causal_conv1d_cuda, fast_causal_conv1d_fn
-    from ...modules.conv.triton import CausalConv1dFunction
-    if cp_context is not None:
-        assert initial_state is None, "Initial state is not supported for CP"
-        assert output_final_state is False, "Output final state is not supported for CP"
-        output = causal_conv1d_cp(
-            x=x,
-            weight=weight,
-            bias=bias,
-            activation=activation,
-            chunk_indices=chunk_indices,
-            cp_context=cp_context,
-        )
-        return output, None
-    if backend == 'triton':
-        y, final_state = CausalConv1dFunction.apply(
-            x,
-            weight,
-            bias,
-            residual,
-            initial_state,
-            output_final_state,
-            activation,
-            cu_seqlens,
-            cu_seqlens_cpu,
-            chunk_indices,
-        )
-        return y, final_state
-    elif backend == 'mix':
-        seq_idx = kwargs.get('seq_idx')
-        return fast_causal_conv1d_fn(
-            x,
-            weight,
-            bias,
-            residual,
-            initial_state,
-            output_final_state,
-            activation,
-            cu_seqlens,
-            cu_seqlens_cpu=cu_seqlens_cpu,
-            chunk_indices=chunk_indices,
-            seq_idx=seq_idx,
-        )
-    elif backend == 'cuda':
-        return causal_conv1d_cuda(
-            x,
-            weight,
-            bias,
-            residual,
-            initial_state,
-            output_final_state,
-            activation,
-            cu_seqlens,
-            cu_seqlens_cpu=cu_seqlens_cpu,
-            **kwargs,
-        )
-    else:
-        raise ValueError(f"Unsupported backend: {backend}")

build/torch-cuda/modules/conv/cuda/ops.py DELETED Viewed

@@ -1,233 +0,0 @@
-# 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
-"""CUDA-based mixed-mode implementation for causal convolution."""
-import torch
-from einops import rearrange
-from ....modules.conv.triton import causal_conv1d_update_states
-from ....ops.utils import prepare_sequence_ids
-from ....utils import input_guard
-try:
-    from causal_conv1d.cpp_functions import causal_conv1d_bwd_function
-except ImportError:
-    causal_conv1d_bwd_function = None
-try:
-    from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda
-except ImportError:
-    causal_conv1d_fn_cuda = None
-class FastCausalConv1dFn(torch.autograd.Function):
-    """
-    Mixed-mode (Mix) Causal Convolution Implementation - Combining Triton Forward and CUDA Backward Propagation
-    This class implements forward propagation using FLA's Triton kernel, while using the optimized
-    implementation from TriDao's causal_conv1d CUDA package for backward propagation.
-    This hybrid strategy combines the advantages of both technologies:
-    - Forward: Uses FLA's Triton implementation, optimized for the FLA framework
-    - Backward: Uses TriDao's causal_conv1d_bwd_function CUDA implementation for faster speed
-    Performance Benefits:
-    - CUDA backward implementation is typically faster than the Triton version, reducing training time
-    - Maintains the flexibility and compatibility of forward propagation
-    Note:
-    - Input/Output format is (batch, seqlen, dim)
-    - Backward propagation requires causal_conv1d package: pip install causal-conv1d
-    - Supports SILU/Swish activation functions
-    - Current limitations (not yet supported):
-        * output_final_state must be False
-        * initial_states must be None
-        * residual must be None
-    """
-    @staticmethod
-    @input_guard(no_guard_contiguous=["x"])
-    def forward(
-        ctx,
-        x,
-        weight,
-        bias=None,
-        residual: torch.Tensor | None = None,
-        initial_states=None,
-        output_final_state=False,
-        activation=None,
-        cu_seqlens: torch.LongTensor | None = None,
-        cu_seqlens_cpu: torch.LongTensor | None = None,
-        chunk_indices: torch.LongTensor | None = None,
-        seq_idx: torch.LongTensor | None = None,
-    ):
-        if activation not in [None, "silu", "swish"]:
-            raise NotImplementedError("activation must be None, silu, or swish")
-        assert output_final_state is False, "output_final_state must be False for FastCausalConv1dFn"
-        assert initial_states is None, "initial_states must be None for FastCausalConv1dFn"
-        assert residual is None, "residual must be None for FastCausalConv1dFn"
-        bias = bias.contiguous() if bias is not None else None
-        if cu_seqlens is not None and seq_idx is None:
-            seq_idx = prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(
-                torch.int32).unsqueeze(0)
-        seq_idx = seq_idx.contiguous() if seq_idx is not None else None
-        # Import here to avoid circular dependency
-        from ....modules.conv.triton.ops import causal_conv1d_fwd
-        ctx.activation = activation in ["silu", "swish"]
-        out, _ = causal_conv1d_fwd(
-            x=x,
-            weight=weight,
-            bias=bias,
-            residual=None,
-            initial_state=None,
-            output_final_state=output_final_state,
-            activation=activation,
-            cu_seqlens=cu_seqlens,
-            cu_seqlens_cpu=cu_seqlens_cpu,
-            chunk_indices=chunk_indices,
-        )
-        ctx.save_for_backward(x, weight, bias, seq_idx, initial_states)
-        ctx.return_final_states = output_final_state
-        ctx.return_dinitial_states = (
-            initial_states is not None and initial_states.requires_grad
-        )
-        return out, None
-    @staticmethod
-    @input_guard
-    def backward(ctx, dout, *args):
-        x, weight, bias, seq_idx, initial_states = ctx.saved_tensors
-        dx = torch.empty_like(x, memory_format=torch.contiguous_format)
-        x = rearrange(x, 'b t d -> b d t')
-        dx = rearrange(dx, 'b t d -> b d t')
-        dout = rearrange(dout, 'b t d -> b d t')
-        dfinal_states = args[0] if ctx.return_final_states else None
-        if dout.stride(2) != 1 and dout.stride(1) != 1:
-            dout = dout.contiguous()
-        # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the
-        # backward of conv1d with the backward of chunk).
-        # Here we just pass in None and dx will be allocated in the C++ code.
-        dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function(
-            x,
-            weight,
-            bias,
-            dout,
-            seq_idx,
-            initial_states,
-            dfinal_states,
-            dx,
-            ctx.return_dinitial_states,
-            ctx.activation,
-        )
-        dx = rearrange(dx, 'b d t -> b t d')
-        return (
-            dx,
-            dweight,
-            dbias if bias is not None else None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
-def fast_causal_conv1d_fn(
-    x: torch.Tensor,
-    weight: torch.Tensor | None = None,
-    bias: torch.Tensor | None = None,
-    residual: torch.Tensor | None = None,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool | None = False,
-    activation: str | None = None,
-    cu_seqlens: torch.Tensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    chunk_indices: torch.LongTensor | None = None,
-    seq_idx: torch.LongTensor | None = None,
-):
-    """
-    x: (batch, seqlen, dim)
-    weight: (dim, width)
-    bias: (dim,)
-    seq_idx: (batch, seqlen)
-    initial_states: (batch, dim, width - 1)
-    final_states_out: (batch, dim, width - 1), to be written to
-    activation: either None or "silu" or "swish"
-    out: (batch, seqlen, dim)
-    """
-    assert causal_conv1d_bwd_function is not None, "causal_conv1d_bwd_function is not available"
-    return FastCausalConv1dFn.apply(
-        x,
-        weight,
-        bias,
-        residual,
-        initial_state,
-        output_final_state,
-        activation,
-        cu_seqlens,
-        cu_seqlens_cpu,
-        chunk_indices,
-        seq_idx,
-    )
-def causal_conv1d_cuda(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor | None = None,
-    residual: torch.Tensor | None = None,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool | None = False,
-    activation: str | None = None,
-    cu_seqlens: torch.Tensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    **kwargs,
-):
-    assert causal_conv1d_fn_cuda is not None, "causal_conv1d_fn_cuda is not available"
-    seq_idx = kwargs.get('seq_idx')
-    if cu_seqlens is not None or seq_idx is not None:
-        assert initial_state is None, "For CUDA backend, initial_state must be None if cu_seqlens or seq_idx is provided"
-    W = weight.shape[-1]
-    if x.stride(-1) != 1:
-        x = x.contiguous()
-    x_conv1d = rearrange(x, 'b t d -> b d t')
-    if cu_seqlens is not None and seq_idx is None:
-        seq_idx = prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(torch.int32).unsqueeze(0)
-    y = causal_conv1d_fn_cuda(
-        x=x_conv1d,
-        weight=weight,
-        bias=bias,
-        activation=activation,
-        seq_idx=seq_idx,
-        initial_states=None,
-        return_final_states=False,
-    )
-    y = rearrange(y, 'b d t -> b t d')
-    if output_final_state:
-        final_state = causal_conv1d_update_states(
-            x=x,
-            state_len=W,
-            initial_state=initial_state,
-            cu_seqlens=cu_seqlens,
-        )
-    else:
-        final_state = None
-    if residual is not None:
-        y.add_(residual)
-    return y, final_state

build/torch-cuda/modules/conv/short_conv.py DELETED Viewed

@@ -1,250 +0,0 @@
-# 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
-"""Short convolution implementation for efficient causal convolutions."""
-import warnings
-import torch
-import torch.nn as nn
-from einops import rearrange
-try:
-    from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda
-    from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda
-except ImportError:
-    causal_conv1d_fn_cuda = None
-    causal_conv1d_update_cuda = None
-class ShortConvolution(nn.Conv1d):
-    """Short convolution layer for efficient causal convolution operations.
-    This class implements a depthwise 1D convolution with causal padding,
-    designed for efficient sequence processing. It supports multiple backends (Triton/CUDA)
-    and optional activation functions.
-    Args:
-        hidden_size (int): Number of input/output channels (must be equal for depthwise conv)
-        kernel_size (int): Size of the convolution kernel
-        bias (bool, optional): Whether to include learnable bias. Defaults to False.
-        activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'.
-        backend (Optional[str], optional): Backend implementation ('triton' or 'cuda'). Defaults to 'triton'.
-        device (Optional[torch.device], optional): Device to place the layer on. Defaults to None.
-        dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None.
-        **kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility)
-    Attributes:
-        hidden_size (int): Number of channels
-        activation (Optional[str]): Selected activation function
-        backend (str): Actual backend being used (may differ from input due to availability)
-    Note:
-        - Uses depthwise convolution (groups=hidden_size) for efficiency
-        - Applies causal padding (kernel_size-1) to ensure no future information leakage
-        - Falls back to Triton backend if CUDA backend is unavailable
-    """
-    def __init__(
-        self,
-        hidden_size: int,
-        kernel_size: int,
-        bias: bool = False,
-        activation: str | None = 'silu',
-        backend: str | None = 'triton',
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-        **kwargs,
-    ):
-        super().__init__(
-            in_channels=hidden_size,
-            out_channels=hidden_size,
-            kernel_size=kernel_size,
-            groups=hidden_size,
-            bias=bias,
-            padding=kernel_size - 1,
-            device=device,
-            dtype=dtype,
-        )
-        self.hidden_size = hidden_size
-        self.activation = None
-        if activation is not None:
-            assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet."
-            self.activation = activation
-        if 'use_fast_conv1d' in kwargs:
-            warnings.warn(
-                "The `use_fast_conv1d` parameter is deprecated and will be ignored. "
-                "Please use the `backend` parameter instead.",
-            )
-        import os
-        self.backend = os.environ.get('FLA_CONV_BACKEND', backend)
-        if backend not in ['cuda', 'triton']:
-            raise ValueError(f"Invalid backend: {backend}, must be one of ['cuda', 'triton']")
-        if backend == 'cuda':
-            if causal_conv1d_fn_cuda is None:
-                warnings.warn(
-                    "The `backend` parameter is set to `cuda`, but `causal_conv1d_fn` is not available. "
-                    "Switching to the Triton implementation instead. "
-                    "Consider installing `causal_conv1d` to enable the CUDA backend.",
-                )
-                self.backend = 'triton'
-    def extra_repr(self):
-        s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}'
-             ', stride={stride}')
-        if self.padding != (0,) * len(self.padding):
-            s += ', padding={padding}'
-        if self.dilation != (1,) * len(self.dilation):
-            s += ', dilation={dilation}'
-        if self.output_padding != (0,) * len(self.output_padding):
-            s += ', output_padding={output_padding}'
-        if self.groups != 1:
-            s += ', groups={groups}'
-        if self.bias is None:
-            s += ', bias=False'
-        if self.padding_mode != 'zeros':
-            s += ', padding_mode={padding_mode}'
-        if self.activation is not None:
-            s += ', activation={activation}'
-        s += f', backend={self.backend}'
-        return s.format(**self.__dict__)
-    def forward(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None = None,
-        mask: torch.Tensor | None = None,
-        cache: torch.Tensor | None = None,
-        output_final_state: bool = False,
-        cu_seqlens: torch.LongTensor | None = None,
-        chunk_indices: torch.LongTensor | None = None,
-        **kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """
-        Args:
-            x (`torch.Tensor`):
-                Tensor of shape `[B, T, D]`. `B` must be 1 if `cu_seqlens` is provided.
-            residual (`Optional[torch.Tensor]`):
-                Residual tensor of shape `[B, T, D]`. Default: `None`.
-            mask (`Optional[torch.Tensor]`):
-                Attention mask dealing with padded positions.
-            cache (`Optional[torch.Tensor]`):
-                Previous cache tensor of shape `[N, D, W]`, where `W` is the kernel size.
-                If provided, the cache is updated **inplace**.
-            output_final_state (Optional[bool]):
-                Whether to output the final state of shape `[N, D, W]`. Default: `False`.
-            cu_seqlens (Optional[torch.LongTensor]):
-                Cumulative sequence lengths for each batch. Used for varlen. Default: `None`.
-                Shape: [B+1]
-            chunk_indices (Optional[torch.LongTensor]):
-                Chunk indices for variable-length sequences. Default: `None`.
-        Returns:
-            Tensor of shape `[B, T, D]`.
-        """
-        # Import here to avoid circular dependency
-        from ...modules.conv.causal_conv1d import causal_conv1d
-        B, T, *_ = x.shape
-        N = B if cu_seqlens is None else len(cu_seqlens) - 1
-        if mask is not None:
-            if cu_seqlens is not None:
-                raise ValueError("`mask` and `cu_seqlens` cannot be provided at the same time")
-            x = x.mul_(mask.unsqueeze(-1))
-        # in decoding phase, the cache (if provided) is updated inplace
-        if B * T == N:
-            y, cache = self.step(
-                x=x,
-                residual=residual,
-                cache=cache,
-                output_final_state=output_final_state,
-                cu_seqlens=cu_seqlens,
-            )
-            return y, cache
-        # cuda backend do not support:
-        # 1. both `cu_seqlens` and `cache` being provided
-        # 2. both `cu_seqlens` and `output_final_state` being provided
-        # and other small issues
-        # to simplify the implementation, we just switch to triton backend
-        if self.backend == 'cuda' and cache is not None:
-            warnings.warn(
-                "The CUDA backend does not support both `cu_seqlens` and `cache` being provided, "
-                "or both `cu_seqlens` and `output_final_state` being provided. "
-                "Switching to the Triton backend instead. ",
-                stacklevel=2,
-            )
-            self.backend = 'triton'
-        return causal_conv1d(
-            x=x,
-            weight=rearrange(self.weight, "d 1 w -> d w"),
-            bias=self.bias,
-            residual=residual,
-            initial_state=cache,
-            output_final_state=output_final_state,
-            activation=self.activation,
-            backend=self.backend,
-            cu_seqlens=cu_seqlens,
-            chunk_indices=chunk_indices,
-            **kwargs,
-        )
-    def step(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None,
-        cache: torch.Tensor | None,
-        output_final_state: bool = False,
-        cu_seqlens: torch.LongTensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        from ...modules.conv.triton.ops import causal_conv1d_update
-        B, _, D, W = *x.shape, self.kernel_size[0]
-        N = B if cu_seqlens is None else len(cu_seqlens) - 1
-        # Always initialise cache when None so the Triton kernel never
-        # receives a None tensor. Return value still respects output_final_state
-        # to maintain consistency with the non-step path in forward().
-        if cache is None:
-            cache = x.new_zeros(N, D, W)
-        # NOTE: we follow the fast mode that updates the cache in-place
-        if self.backend == 'triton':
-            y, cache = causal_conv1d_update(
-                x=x,
-                cache=cache,
-                residual=residual,
-                weight=rearrange(self.weight, "d 1 w -> d w"),
-                bias=self.bias,
-                activation=self.activation,
-            )
-            return y, (cache if output_final_state else None)
-        shape = x.shape
-        x = x.squeeze(0) if cu_seqlens is not None else x.squeeze(1)
-        # equivalent to:
-        # cache.copy_(cache.roll(shifts=-1, dims=-1))
-        # cache[:, :, -1] = x
-        # y = torch.sum(cache * rearrange(self.weight, "d 1 w -> d w"), dim=-1)
-        y = causal_conv1d_update_cuda(
-            x=x,
-            conv_state=cache,
-            weight=rearrange(self.weight, "d 1 w -> d w"),
-            bias=self.bias,
-            activation=self.activation,
-        )
-        y = y.view(shape)
-        if residual is not None:
-            y.add_(residual)
-        return y, (cache if output_final_state else None)
-    @property
-    def state_size(self) -> int:
-        return self.hidden_size * self.kernel_size

build/torch-cuda/modules/conv/triton/ops.py DELETED Viewed

@@ -1,424 +0,0 @@
-# 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
-import torch
-import triton
-from einops import rearrange
-from ....modules.backends import dispatch
-from ....ops.utils import prepare_chunk_indices
-from ....utils import input_guard
-from .kernels import (
-    causal_conv1d_bwd_kernel,
-    causal_conv1d_fwd_kernel,
-    causal_conv1d_states_fwd_kernel,
-    causal_conv1d_update_kernel,
-    compute_dh0_kernel,
-)
-def _has_non_standard_layout(x: torch.Tensor) -> bool:
-    """QKV-style views (stride_t != D) break triton-ascend masked kernels."""
-    if x.dtype not in (torch.float16, torch.bfloat16):
-        return False
-    if x.dim() != 3:
-        return not x.is_contiguous()
-    _, stride_t, stride_d = x.stride()
-    return stride_d == 1 and stride_t != x.shape[-1]
-@dispatch('modules')
-@input_guard(no_guard_contiguous=["x"])
-def causal_conv1d_fwd(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    residual: torch.Tensor,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool = False,
-    activation: str | None = None,
-    cu_seqlens: torch.LongTensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    chunk_indices: torch.LongTensor | None = None,
-    BT: int = 64,
-    layout_fallback: bool = False,
-) -> torch.Tensor:
-    shape = x.shape
-    if x.shape[-1] != weight.shape[0]:
-        x = rearrange(x, 'b t ... -> b t (...)')
-    B, T, D = x.shape[0], x.shape[1], weight.shape[0]
-    W = weight.shape[1]
-    stride_x_n, stride_x_t, stride_x_d = x.stride()
-    BW = triton.next_power_of_2(W)
-    if cu_seqlens is not None and chunk_indices is None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
-    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
-    NB = triton.cdiv(B*T, 1024)
-    y = torch.empty_like(x, memory_format=torch.contiguous_format)
-    def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B)
-    causal_conv1d_fwd_kernel[grid](
-        x=x,
-        y=y,
-        weight=weight,
-        bias=bias,
-        residual=residual,
-        cu_seqlens=cu_seqlens,
-        initial_state=initial_state,
-        chunk_indices=chunk_indices,
-        B=B,
-        T=T,
-        D=D,
-        W=W,
-        BT=BT,
-        BW=BW,
-        NB=NB,
-        stride_x_n=stride_x_n,
-        stride_x_t=stride_x_t,
-        stride_x_d=stride_x_d,
-        ACTIVATION=activation,
-    )
-    final_state = None
-    if output_final_state:
-        final_state = causal_conv1d_update_states(
-            x=x,
-            state_len=W,
-            initial_state=initial_state,
-            cu_seqlens=cu_seqlens,
-        )
-    return y.view(shape), final_state
-@dispatch('modules')
-def compute_dh0_triton(
-    dy: torch.Tensor,
-    y: torch.Tensor | None,
-    weight: torch.Tensor,
-    initial_state: torch.Tensor,
-    activation: str | None,
-    cu_seqlens: torch.Tensor | None,
-) -> torch.Tensor:
-    """
-    Compute dh0 (gradient w.r.t. initial_state) using a separate Triton kernel.
-    This is a workaround for Triton compiler bugs on some architectures (e.g., GB200).
-    """
-    D, W = weight.shape
-    N = initial_state.shape[0]
-    T = dy.shape[1]
-    # Initialize dh0
-    dh0 = torch.zeros_like(initial_state)
-    BD = 32
-    grid = (triton.cdiv(D, BD), N)
-    y_to_pass = y if activation in ('swish', 'silu') else None
-    # dy is [B, T, D], stride_n = T*D, stride_t = D
-    stride_dy_n = dy.stride(0)
-    stride_dy_t = dy.stride(1)
-    compute_dh0_kernel[grid](
-        dy=dy,
-        y=y_to_pass,
-        weight=weight,
-        dh0=dh0,
-        cu_seqlens=cu_seqlens,
-        stride_dy_n=stride_dy_n,
-        stride_dy_t=stride_dy_t,
-        T=T,
-        D=D,
-        W=W,
-        BD=BD,
-    )
-    return dh0
-@dispatch('modules')
-def causal_conv1d_bwd(
-    x: torch.Tensor,
-    dy: torch.Tensor,
-    dht: torch.Tensor,
-    weight: torch.Tensor | None = None,
-    bias: torch.Tensor | None = None,
-    residual: torch.Tensor | None = None,
-    initial_state: torch.Tensor | None = None,
-    activation: str | None = None,
-    cu_seqlens: torch.Tensor | None = None,
-    cu_seqlens_cpu: torch.LongTensor | None = None,
-    chunk_indices: torch.LongTensor | None = None,
-    BT: int = 64,
-    layout_fallback: bool = False,
-):
-    shape = x.shape
-    if x.shape[-1] != weight.shape[0]:
-        x = rearrange(x, 'b t ... -> b t (...)')
-    B, T, D = x.shape
-    W = weight.shape[1] if weight is not None else None
-    stride_x_n, stride_x_t, stride_x_d = x.stride()
-    stride_dy_n, stride_dy_t, stride_dy_d = dy.stride()
-    BW = triton.next_power_of_2(W)
-    if cu_seqlens is not None and chunk_indices is None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
-    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
-    NB = triton.cdiv(B*T, 1024)
-    y = None
-    if activation is not None:
-        y, _ = causal_conv1d_fwd(
-            x=x,
-            weight=weight,
-            bias=bias,
-            residual=None,
-            initial_state=initial_state,
-            activation=None,
-            cu_seqlens=cu_seqlens,
-            cu_seqlens_cpu=cu_seqlens_cpu,
-            output_final_state=False,
-            chunk_indices=chunk_indices,
-        )
-    dx = torch.empty_like(x)
-    dw = weight.new_empty(B*NT, *weight.shape, dtype=torch.float) if weight is not None else None
-    db = bias.new_empty(B*NT, *bias.shape, dtype=torch.float) if bias is not None else None
-    dr = dy if residual is not None else None
-    stride_dx_n, stride_dx_t, stride_dx_d = dx.stride()
-    def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B)
-    causal_conv1d_bwd_kernel[grid](
-        x=x,
-        y=y,
-        weight=weight,
-        initial_state=initial_state,
-        dht=dht,
-        dy=dy,
-        dx=dx,
-        dw=dw,
-        db=db,
-        cu_seqlens=cu_seqlens,
-        chunk_indices=chunk_indices,
-        B=B,
-        T=T,
-        D=D,
-        W=W,
-        BT=BT,
-        BW=BW,
-        NB=NB,
-        stride_x_n=stride_x_n,
-        stride_x_t=stride_x_t,
-        stride_x_d=stride_x_d,
-        stride_dx_n=stride_dx_n,
-        stride_dx_t=stride_dx_t,
-        stride_dx_d=stride_dx_d,
-        stride_dy_n=stride_dy_n,
-        stride_dy_t=stride_dy_t,
-        stride_dy_d=stride_dy_d,
-        ACTIVATION=activation,
-    )
-    if weight is not None:
-        dw = dw.sum(0).to(weight)
-    if bias is not None:
-        db = db.sum(0).to(bias)
-    # Compute dh0 using separate Triton kernel to avoid compiler bugs on some architectures (e.g., GB200)
-    dh0 = None
-    if initial_state is not None:
-        dh0 = compute_dh0_triton(
-            dy=dy,
-            y=y,
-            weight=weight,
-            initial_state=initial_state,
-            activation=activation,
-            cu_seqlens=cu_seqlens,
-        )
-    return dx.view(shape), dw, db, dr, dh0
-@dispatch('modules')
-@input_guard(no_guard_contiguous=["x"])
-def causal_conv1d_update_states(
-    x: torch.Tensor,
-    state_len: int,
-    initial_state: torch.Tensor | None = None,
-    cu_seqlens: torch.Tensor | None = None,
-) -> torch.Tensor:
-    if cu_seqlens is not None:
-        N = len(cu_seqlens) - 1
-        if x.dim() == 2:
-            stride_x_n = 0
-            stride_x_t, stride_x_d = x.stride()
-            T = x.shape[0]
-        else:
-            stride_x_n = x.stride(0)
-            stride_x_t, stride_x_d = x.stride(1), x.stride(2)
-            T = x.shape[1]
-        D = x.shape[-1]
-    else:
-        B, T, D = x.shape
-        N = B
-        stride_x_n, stride_x_t, stride_x_d = x.stride()
-    W = state_len
-    final_state = torch.empty(N, D, W, dtype=x.dtype, device=x.device)
-    BD = min(triton.next_power_of_2(D), 256)
-    BW = triton.next_power_of_2(W)
-    grid = (triton.cdiv(D, BD), N)
-    causal_conv1d_states_fwd_kernel[grid](
-        x=x,
-        initial_state=initial_state,
-        final_state=final_state,
-        cu_seqlens=cu_seqlens,
-        T=T,
-        D=D,
-        W=W,
-        stride_x_n=stride_x_n,
-        stride_x_t=stride_x_t,
-        stride_x_d=stride_x_d,
-        BW=BW,
-        BD=BD,
-    )
-    return final_state
-@dispatch('modules')
-@input_guard(no_guard_contiguous=["x"])
-def causal_conv1d_update(
-    x: torch.Tensor,
-    cache: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    weight: torch.Tensor | None = None,
-    bias: torch.Tensor | None = None,
-    activation: str | None = None,
-) -> torch.Tensor:
-    shape = x.shape
-    if weight is not None and x.shape[-1] != weight.shape[0]:
-        x = rearrange(x, 'b t ... -> b t (...)')
-    D = x.shape[-1]
-    N = x.numel() // D
-    W = weight.shape[1] if weight is not None else None
-    BW = triton.next_power_of_2(W)
-    if x.dim() == 2:
-        # Case: (N, D)
-        stride_x_n = x.stride(0)
-        stride_x_d = x.stride(1)
-    elif x.dim() == 3 and x.shape[0] == 1:
-        # Case: (1, N, D) -> Time=1, Batch=N, Dim=D
-        # Batch 在 dim 1
-        stride_x_n = x.stride(1)
-        stride_x_d = x.stride(2)
-    elif x.dim() == 3:
-        # Case: (N, 1, D) -> Batch=N, Time=1, Dim=D
-        # Batch 在 dim 0
-        stride_x_n = x.stride(0)
-        stride_x_d = x.stride(2)
-    else:
-        # Fallback / Error case
-        raise ValueError(f"Unsupported input shape: {x.shape}")
-    y = torch.empty_like(x, memory_format=torch.contiguous_format)
-    if y.dim() == 2:
-        stride_y_n, stride_y_d = y.stride(0), y.stride(1)
-    elif y.dim() == 3 and y.shape[0] == 1:
-        stride_y_n, stride_y_d = y.stride(1), y.stride(2)
-    elif y.dim() == 3:
-        stride_y_n, stride_y_d = y.stride(0), y.stride(2)
-    def grid(meta): return (triton.cdiv(D, meta['BD']), N)
-    causal_conv1d_update_kernel[grid](
-        x=x,
-        cache=cache,
-        residual=residual,
-        y=y,
-        weight=weight,
-        bias=bias,
-        stride_x_n=stride_x_n,
-        stride_x_d=stride_x_d,
-        stride_y_n=stride_y_n,
-        stride_y_d=stride_y_d,
-        D=D,
-        W=W,
-        BW=BW,
-        ACTIVATION=activation,
-    )
-    return y.view(shape), cache
-class CausalConv1dFunction(torch.autograd.Function):
-    @staticmethod
-    @input_guard(no_guard_contiguous=["x"])
-    def forward(
-        ctx,
-        x: torch.Tensor,
-        weight: torch.Tensor | None = None,
-        bias: torch.Tensor | None = None,
-        residual: torch.Tensor | None = None,
-        initial_state: torch.Tensor | None = None,
-        output_final_state: bool | None = False,
-        activation: str | None = None,
-        cu_seqlens: torch.Tensor | None = None,
-        cu_seqlens_cpu: torch.LongTensor | None = None,
-        chunk_indices: torch.LongTensor | None = None,
-        chunk_size: int = 64,
-    ):
-        BT = chunk_size
-        if cu_seqlens is not None and chunk_indices is None:
-            chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu)
-        ctx.activation = activation
-        ctx.cu_seqlens = cu_seqlens
-        ctx.cu_seqlens_cpu = cu_seqlens_cpu
-        ctx.chunk_indices = chunk_indices
-        ctx.layout_fallback = _has_non_standard_layout(x)
-        ctx.save_for_backward(x, weight, bias, residual, initial_state)
-        y, final_state = causal_conv1d_fwd(
-            x=x,
-            weight=weight,
-            bias=bias,
-            residual=residual,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            activation=activation,
-            cu_seqlens=cu_seqlens,
-            cu_seqlens_cpu=cu_seqlens_cpu,
-            chunk_indices=chunk_indices,
-            BT=BT,
-            layout_fallback=ctx.layout_fallback,
-        )
-        return y, final_state
-    @staticmethod
-    @input_guard(no_guard_contiguous=["dy"])
-    def backward(ctx, dy: torch.Tensor, dht: torch.Tensor | None = None):
-        x, weight, bias, residual, initial_state = ctx.saved_tensors
-        dx, dw, db, dr, dh0 = causal_conv1d_bwd(
-            x=x,
-            dy=dy,
-            dht=dht,
-            weight=weight,
-            bias=bias,
-            residual=residual,
-            initial_state=initial_state,
-            activation=ctx.activation,
-            cu_seqlens=ctx.cu_seqlens,
-            cu_seqlens_cpu=ctx.cu_seqlens_cpu,
-            chunk_indices=ctx.chunk_indices,
-            layout_fallback=ctx.layout_fallback,
-        )
-        return dx, dw, db, dr, dh0, None, None, None, None, None, None

build/torch-cuda/modules/convolution.py DELETED Viewed

@@ -1,42 +0,0 @@
-# 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 ..modules.conv import (
-    ImplicitLongConvolution,
-    LongConvolution,
-    PositionalEmbedding,
-    ShortConvolution,
-    causal_conv1d,
-    fft_conv,
-)
-from ..modules.conv.cp import CausalConv1dFunctionCP, causal_conv1d_cp
-from ..modules.conv.cuda import FastCausalConv1dFn, fast_causal_conv1d_fn
-from ..modules.conv.triton import (
-    CausalConv1dFunction,
-    causal_conv1d_bwd,
-    causal_conv1d_fwd,
-    causal_conv1d_update,
-    causal_conv1d_update_states,
-)
-__all__ = [
-    'CausalConv1dFunction',
-    'CausalConv1dFunctionCP',
-    'FastCausalConv1dFn',
-    'ImplicitLongConvolution',
-    'LongConvolution',
-    'PositionalEmbedding',
-    'ShortConvolution',
-    'causal_conv1d',
-    'causal_conv1d_bwd',
-    'causal_conv1d_cp',
-    'causal_conv1d_fwd',
-    'causal_conv1d_update',
-    'causal_conv1d_update_states',
-    'fast_causal_conv1d_fn',
-    'fft_conv',
-]

build/torch-cuda/modules/feature_map.py CHANGED Viewed

@@ -121,8 +121,7 @@ class DPFPFeatureMap(nn.Module):
         self.nu = nu
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        relu = x.relu()
-        x = torch.cat([relu, -relu], dim=-1)
         x_rolled = torch.cat([x.roll(shifts=j, dims=-1) for j in range(1, self.nu+1)], dim=-1)
         x_repeat = torch.cat([x] * self.nu, dim=-1)
         return x_repeat * x_rolled

         self.nu = nu
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.cat([x.relu(), (-x).relu()], dim=-1)
         x_rolled = torch.cat([x.roll(shifts=j, dims=-1) for j in range(1, self.nu+1)], dim=-1)
         x_repeat = torch.cat([x] * self.nu, dim=-1)
         return x_repeat * x_rolled

build/torch-cuda/modules/fused_bitlinear.py CHANGED Viewed

@@ -5,491 +5,42 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-# Implementations of BitLinear layer with fused LayerNorm and quantized Linear layer.
 # [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764)
 # [Scalable MatMul-free Language Modeling](https://arxiv.org/abs/2406.02528)
 #
-# Code adapted from https://github.com/ridgerchu/matmulfreellm/
 from __future__ import annotations
-import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import triton
-import triton.language as tl
-from ..modules.layernorm import RMSNorm
-from ..utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard, require_version
-NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32]
-def activation_quant(x):
-    """
-    Per-token quantization to 8 bits. No grouping is needed for quantization.
-    Args:
-        x: An activation tensor with shape [n, d].
-    Returns:
-        A quantized activation tensor with shape [n, d].
-    """
-    # Compute the scale factor
     scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
-    # Quantize and then de-quantize the tensor
-    y = (x * scale).round().clamp_(-128, 127) / scale
     return y
-def weight_quant(w):
-    """
-    Per-tensor quantization to 1.58 bits. No grouping is needed for quantization.
-    Args:
-        w: A weight tensor with shape [d, k].
-    Returns:
-        A quantized weight tensor with shape [d, k].
-    """
-    # Compute the scale factor
-    scale = 1.0 / w.abs().mean().clamp_(min=1e-5)
-    # Quantize and then de-quantize the tensor
-    u = (w * scale).round().clamp_(-1, 1) / scale
-    return u
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=num_warps)
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS"],
-    **autotune_cache_kwargs,
-)
-@triton.jit
-def layer_norm_fwd_kernel_quant(
-    X,  # pointer to the input
-    Y,  # pointer to the output
-    W,  # pointer to the weights
-    B,  # pointer to the biases
-    RESIDUAL,  # pointer to the residual
-    RESIDUAL_OUT,  # pointer to the residual
-    Mean,  # pointer to the mean
-    Rstd,  # pointer to the 1/std
-    stride_x_row,  # how much to increase the pointer when moving by 1 row
-    stride_y_row,
-    stride_res_row,
-    stride_res_out_row,
-    N,  # number of columns in X
-    eps,  # epsilon to avoid division by zero
-    IS_RMS_NORM: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    HAS_RESIDUAL: tl.constexpr,
-    STORE_RESIDUAL_OUT: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-):
-    # Map the program id to the row of X and Y it should compute.
-    row = tl.program_id(0)
-    X += row * stride_x_row
-    Y += row * stride_y_row
-    if HAS_RESIDUAL:
-        RESIDUAL += row * stride_res_row
-    if STORE_RESIDUAL_OUT:
-        RESIDUAL_OUT += row * stride_res_out_row
-    # Compute mean and variance
-    cols = tl.arange(0, BLOCK_N)
-    x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
-    if HAS_RESIDUAL:
-        residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32)
-        x += residual
-    if STORE_RESIDUAL_OUT:
-        tl.store(RESIDUAL_OUT + cols, x, mask=cols < N)
-    if not IS_RMS_NORM:
-        mean = tl.sum(x, axis=0) / N
-        tl.store(Mean + row, mean)
-        xbar = tl.where(cols < N, x - mean, 0.0)
-        var = tl.sum(xbar * xbar, axis=0) / N
-    else:
-        xbar = tl.where(cols < N, x, 0.0)
-        var = tl.sum(xbar * xbar, axis=0) / N
-    rstd = 1 / tl.sqrt(var + eps)
-    tl.store(Rstd + row, rstd)
-    # Normalize and apply linear transformation
-    mask = cols < N
-    if HAS_WEIGHT:
-        w = tl.load(W + cols, mask=mask).to(tl.float32)
-    if HAS_BIAS:
-        b = tl.load(B + cols, mask=mask).to(tl.float32)
-    x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
-    y = x_hat * w if HAS_WEIGHT else x_hat
-    if HAS_BIAS:
-        y = y + b
-    # Aply quantization to the output
-    scale = 127.0 / tl.maximum(tl.max(tl.abs(y), 0), 1e-5)
-    # Quantize and then de-quantize the tensor
-    y = tl.extra.cuda.libdevice.round(y * scale)
-    y = tl.maximum(tl.minimum(y, 127), -128) / scale
-    # Write output
-    tl.store(Y + cols, y, mask=mask)
-def layer_norm_fwd_quant(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    eps: float,
-    residual: torch.Tensor = None,
-    out_dtype: torch.dtype = None,
-    residual_dtype: torch.dtype = None,
-    is_rms_norm: bool = False,
-):
-    if residual is not None:
-        residual_dtype = residual.dtype
-    M, N = x.shape
-    # allocate output
-    y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
-    if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype):
-        residual_out = torch.empty(M, N, device=x.device, dtype=residual_dtype)
-    else:
-        residual_out = None
-    mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None
-    rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
-    # Less than 64KB per feature: enqueue fused kernel
-    MAX_FUSED_SIZE = 65536 // x.element_size()
-    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
-    if N > BLOCK_N:
-        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-    # heuristics for number of warps
-    layer_norm_fwd_kernel_quant[(M,)](
-        x,
-        y,
-        weight,
-        bias,
-        residual,
-        residual_out,
-        mean,
-        rstd,
-        x.stride(0),
-        y.stride(0),
-        residual.stride(0) if residual is not None else 0,
-        residual_out.stride(0) if residual_out is not None else 0,
-        N,
-        eps,
-        is_rms_norm,
-        BLOCK_N,
-        residual is not None,
-        residual_out is not None,
-        weight is not None,
-        bias is not None,
-    )
-    # residual_out is None if residual is None and residual_dtype == input_dtype
-    return y, mean, rstd, residual_out if residual_out is not None else x
-@triton.heuristics({
-    "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None,
-})
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=num_warps)
-        for num_warps in NUM_WARPS_AUTOTUNE
-    ],
-    key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS"],
-    **autotune_cache_kwargs,
-)
-@triton.jit
-def layer_norm_bwd_kernel(
-    X,  # pointer to the input
-    W,  # pointer to the weights
-    B,  # pointer to the biases
-    Y,  # pointer to the output to be recomputed
-    DY,  # pointer to the output gradient
-    DX,  # pointer to the input gradient
-    DW,  # pointer to the partial sum of weights gradient
-    DB,  # pointer to the partial sum of biases gradient
-    DRESIDUAL,
-    DRESIDUAL_IN,
-    Mean,  # pointer to the mean
-    Rstd,  # pointer to the 1/std
-    stride_x_row,  # how much to increase the pointer when moving by 1 row
-    stride_y_row,
-    stride_dy_row,
-    stride_dx_row,
-    stride_dres_row,
-    stride_dres_in_row,
-    M,  # number of rows in X
-    N,  # number of columns in X
-    eps,  # epsilon to avoid division by zero
-    rows_per_program,
-    IS_RMS_NORM: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    HAS_DRESIDUAL: tl.constexpr,
-    STORE_DRESIDUAL: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    RECOMPUTE_OUTPUT: tl.constexpr,
-):
-    # Map the program id to the elements of X, DX, and DY it should compute.
-    row_block_id = tl.program_id(0)
-    row_start = row_block_id * rows_per_program
-    cols = tl.arange(0, BLOCK_N)
-    mask = cols < N
-    X += row_start * stride_x_row
-    if HAS_DRESIDUAL:
-        DRESIDUAL += row_start * stride_dres_row
-    if STORE_DRESIDUAL:
-        DRESIDUAL_IN += row_start * stride_dres_in_row
-    DY += row_start * stride_dy_row
-    DX += row_start * stride_dx_row
-    if RECOMPUTE_OUTPUT:
-        Y += row_start * stride_y_row
-    if HAS_WEIGHT:
-        w = tl.load(W + cols, mask=mask).to(tl.float32)
-        dw = tl.zeros((BLOCK_N,), dtype=tl.float32)
-    if RECOMPUTE_OUTPUT and HAS_BIAS:
-        b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32)
-    if HAS_BIAS:
-        db = tl.zeros((BLOCK_N,), dtype=tl.float32)
-    row_end = min((row_block_id + 1) * rows_per_program, M)
-    for row in range(row_start, row_end):
-        # Load data to SRAM
-        x = tl.load(X + cols, mask=mask, other=0).to(tl.float32)
-        dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32)
-        if not IS_RMS_NORM:
-            mean = tl.load(Mean + row)
-        rstd = tl.load(Rstd + row)
-        # Compute dx
-        xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
-        xhat = tl.where(mask, xhat, 0.0)
-        if RECOMPUTE_OUTPUT:
-            y = xhat * w if HAS_WEIGHT else xhat
-            if HAS_BIAS:
-                y = y + b
-            # Aply quantization to the output
-            scale = 127.0 / tl.maximum(tl.max(tl.abs(y), 0), 1e-5)
-            # Quantize and then de-quantize the tensor
-            y = tl.extra.cuda.libdevice.round(y * scale)
-            y = tl.maximum(tl.minimum(y, 127), -128) / scale
-            tl.store(Y + cols, y, mask=mask)
-        wdy = dy
-        if HAS_WEIGHT:
-            wdy = dy * w
-            dw += dy * xhat
-        if HAS_BIAS:
-            db += dy
-        if not IS_RMS_NORM:
-            c1 = tl.sum(xhat * wdy, axis=0) / N
-            c2 = tl.sum(wdy, axis=0) / N
-            dx = (wdy - (xhat * c1 + c2)) * rstd
-        else:
-            c1 = tl.sum(xhat * wdy, axis=0) / N
-            dx = (wdy - xhat * c1) * rstd
-        if HAS_DRESIDUAL:
-            dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32)
-            dx += dres
-        # Write dx
-        if STORE_DRESIDUAL:
-            tl.store(DRESIDUAL_IN + cols, dx, mask=mask)
-        tl.store(DX + cols, dx, mask=mask)
-        X += stride_x_row
-        if HAS_DRESIDUAL:
-            DRESIDUAL += stride_dres_row
-        if STORE_DRESIDUAL:
-            DRESIDUAL_IN += stride_dres_in_row
-        if RECOMPUTE_OUTPUT:
-            Y += stride_y_row
-        DY += stride_dy_row
-        DX += stride_dx_row
-    if HAS_WEIGHT:
-        tl.store(DW + row_block_id * N + cols, dw, mask=mask)
-    if HAS_BIAS:
-        tl.store(DB + row_block_id * N + cols, db, mask=mask)
-def layer_norm_bwd(
-    dy: torch.Tensor,
     x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    eps: float,
-    mean: torch.Tensor,
-    rstd: torch.Tensor,
-    dresidual: torch.Tensor = None,
-    has_residual: bool = False,
     is_rms_norm: bool = False,
-    x_dtype: torch.dtype = None,
-    recompute_output: bool = False,
-):
-    M, N = x.shape
-    # allocate output
-    dx = torch.empty_like(x) if x_dtype is None else torch.empty(M, N, dtype=x_dtype, device=x.device)
-    dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None
-    y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None
-    # Less than 64KB per feature: enqueue fused kernel
-    MAX_FUSED_SIZE = 65536 // x.element_size()
-    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
-    if N > BLOCK_N:
-        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-    sm_count = get_multiprocessor_count(x.device.index)
-    _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) if weight is not None else None
-    _db = torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) if bias is not None else None
-    rows_per_program = math.ceil(M / sm_count)
-    grid = (sm_count,)
-    layer_norm_bwd_kernel[grid](
-        x,
-        weight,
-        bias,
-        y,
-        dy,
-        dx,
-        _dw,
-        _db,
-        dresidual,
-        dresidual_in,
-        mean,
-        rstd,
-        x.stride(0),
-        0 if not recompute_output else y.stride(0),
-        dy.stride(0),
-        dx.stride(0),
-        dresidual.stride(0) if dresidual is not None else 0,
-        dresidual_in.stride(0) if dresidual_in is not None else 0,
-        M,
-        N,
-        eps,
-        rows_per_program,
-        is_rms_norm,
-        BLOCK_N,
-        dresidual is not None,
-        dresidual_in is not None,
-        weight is not None,
-        bias is not None,
-    )
-    dw = _dw.sum(0).to(weight.dtype) if weight is not None else None
-    db = _db.sum(0).to(bias.dtype) if bias is not None else None
-    # Don't need to compute dresidual_in separately in this case
-    if has_residual and dx.dtype == x.dtype:
-        dresidual_in = dx
-    return (dx, dw, db, dresidual_in) if not recompute_output else (dx, dw, db, dresidual_in, y)
-class LayerNormLinearQuantFn(torch.autograd.Function):
-    @staticmethod
-    @input_guard
-    def forward(
-        ctx,
-        x,
-        norm_weight,
-        norm_bias,
-        linear_weight,
-        linear_bias,
-        residual=None,
-        eps=1e-6,
-        prenorm=False,
-        residual_in_fp32=False,
-        is_rms_norm=False,
-    ):
-        x_shape_og = x.shape
-        # reshape input data into 2D tensor
-        x = x.reshape(-1, x.shape[-1])
-        if residual is not None:
-            assert residual.shape == x_shape_og
-            residual = residual.reshape(-1, residual.shape[-1])
-        residual_dtype = residual.dtype if residual is not None else (torch.float32 if residual_in_fp32 else None)
-        y, mean, rstd, residual_out = layer_norm_fwd_quant(
-            x,
-            norm_weight,
-            norm_bias,
-            eps,
-            residual,
-            out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype(),
-            residual_dtype=residual_dtype,
-            is_rms_norm=is_rms_norm,
-        )
-        y = y.reshape(x_shape_og)
-        dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype
-        linear_weight = weight_quant(linear_weight).to(dtype)
-        linear_bias = linear_bias.to(dtype) if linear_bias is not None else None
-        out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias)
-        # We don't store y, will be recomputed in the backward pass to save memory
-        ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd)
-        ctx.x_shape_og = x_shape_og
-        ctx.eps = eps
-        ctx.is_rms_norm = is_rms_norm
-        ctx.has_residual = residual is not None
-        ctx.prenorm = prenorm
-        ctx.x_dtype = x.dtype
-        ctx.linear_bias_is_none = linear_bias is None
-        return out if not prenorm else (out, residual_out.reshape(x_shape_og))
-    @staticmethod
-    @input_guard
-    def backward(ctx, dout, *args):
-        x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors
-        dout = dout.reshape(-1, dout.shape[-1])
-        dy = F.linear(dout, linear_weight.t())
-        dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
-        assert dy.shape == x.shape
-        if ctx.prenorm:
-            dresidual = args[0]
-            dresidual = dresidual.reshape(-1, dresidual.shape[-1])
-            assert dresidual.shape == x.shape
-        else:
-            dresidual = None
-        dx, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_bwd(
-            dy,
-            x,
-            norm_weight,
-            norm_bias,
-            ctx.eps,
-            mean,
-            rstd,
-            dresidual,
-            ctx.has_residual,
-            ctx.is_rms_norm,
-            x_dtype=ctx.x_dtype,
-            recompute_output=True,
-        )
-        dlinear_weight = torch.einsum("bo,bi->oi", dout, y)
-        return (
-            dx.reshape(ctx.x_shape_og),
-            dnorm_weight,
-            dnorm_bias,
-            dlinear_weight,
-            dlinear_bias,
-            dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None,
-            None,
-            None,
-            None,
-            None,
-        )
-def layer_norm_linear_quant_fn(
-    x,
-    norm_weight,
-    norm_bias,
-    linear_weight,
-    linear_bias,
-    residual=None,
-    eps=1e-6,
-    prenorm=False,
-    residual_in_fp32=False,
-    is_rms_norm=False,
-):
-    return LayerNormLinearQuantFn.apply(
         x,
         norm_weight,
         norm_bias,
@@ -505,16 +56,16 @@ def layer_norm_linear_quant_fn(
 def rms_norm_linear_quant(
     x: torch.Tensor,
-    norm_weight: torch.Tensor,
-    norm_bias: torch.Tensor,
     linear_weight: torch.Tensor,
-    linear_bias: torch.Tensor,
-    residual: torch.Tensor = None,
     eps: float = 1e-5,
     prenorm: bool = False,
     residual_in_fp32: bool = False,
-):
-    return layer_norm_linear_quant_fn(
         x=x,
         norm_weight=norm_weight,
         norm_bias=norm_bias,
@@ -528,111 +79,77 @@ def rms_norm_linear_quant(
     )
-@require_version("triton>=3.0", "Triton >= 3.0 is required to do online quantization.")
-def bit_linear(x, weight, bias=None, norm_weight=None, norm_bias=None, eps=1e-8):
-    """
-    A functional version of BitLinear that applies quantization to activations and weights.
-    Args:
-        x: Input tensor with shape [n, d].
-        weight: Weight tensor with shape [out_features, in_features].
-        bias: Bias tensor with shape [out_features] (optional).
-        norm_weight: Weight tensor for RMS normalization with shape [in_features].
-        norm_bias: Bias tensor for RMS normalization with shape [in_features].
-        eps: A small constant for numerical stability in normalization.
-    Returns:
-        Output tensor with shape [n, out_features].
-    """
-    return layer_norm_linear_quant_fn(
-        x,
-        norm_weight,
-        norm_bias,
-        weight,
-        bias,
         is_rms_norm=True,
     )
 class BitLinear(nn.Linear):
     """
-    A custom linear layer that applies quantization on both activations and weights.
-    This is primarily for training; kernel optimization is needed for efficiency in deployment.
-    """
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        bias: bool = False,
-        norm_eps: float = 1e-8,
-    ):
-        """
-        Initializes the BitLinear layer.
-        Args:
-            in_features: Size of each input sample.
-            out_features: Size of each output sample.
-            bias: If set to False, the layer will not learn an additive bias. Default: True.
-        """
-        # Initialize the superclass nn.Linear with the given parameters
-        super().__init__(in_features, out_features, bias=bias)
-        self.norm = RMSNorm(in_features, eps=norm_eps, dtype=torch.float32)
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({super().extra_repr()}, norm_eps={self.norm.eps})"
-    def forward(self, x):
-        """
-        Overrides the forward pass to include quantization.
-        Args:
-            x: An input tensor with shape [n, d].
-        Returns:
-            An output tensor with shape [n, d].
-        """
-        # Weight tensor
         w = self.weight
-        # Apply RMS normalization to the input
-        x_norm = self.norm(x)
-        # Apply quantization to both activations and weights
-        # Uses Straight-Through Estimator (STE) trick with .detach() for gradient flow
-        x_quant = x_norm + (activation_quant(x_norm) - x_norm).detach()
-        w_quant = w + (weight_quant(w) - w).detach()
-        # Perform linear operation with quantized values
-        y = F.linear(x_quant, w_quant)
         return y
 class FusedBitLinear(BitLinear):
-    """
-    A custom linear layer that applies quantization on both activations and weights.
-    This is primarily for training; kernel optimization is needed for efficiency in deployment.
-    """
-    def __init__(self, in_features, out_features, bias=False):
-        """
-        Initializes the BitLinear layer.
-        Args:
-            in_features: Size of each input sample.
-            out_features: Size of each output sample.
-            bias: If set to False, the layer will not learn an additive bias. Default: True.
-        """
-        # Initialize the superclass nn.Linear with the given parameters
-        super().__init__(in_features, out_features, bias=bias)
-    def forward(self, x):
-        return layer_norm_linear_quant_fn(
-            x,
-            self.norm.weight,
-            self.norm.bias,
-            self.weight,
-            self.bias,
             is_rms_norm=True,
         )

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+# BitLinear with fused normalization and quantized linear projection.
 # [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764)
 # [Scalable MatMul-free Language Modeling](https://arxiv.org/abs/2406.02528)
 #
+# adapted from https://github.com/ridgerchu/matmulfreellm/
 from __future__ import annotations
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from ..modules.norm import RMSNorm
+from ..modules.norm.layernorm_quant import LayerNormLinearQuantFunction, weight_quant
+def activation_quant(x: torch.Tensor) -> torch.Tensor:
+    """Quantize activations to 8 bits per token and return dequantized values."""
     scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
+    y = (x * scale).round().clamp_(min=-128, max=127) / scale
     return y
+def layer_norm_linear_quant(
     x: torch.Tensor,
+    norm_weight: torch.Tensor | None,
+    norm_bias: torch.Tensor | None,
+    linear_weight: torch.Tensor,
+    linear_bias: torch.Tensor | None,
+    residual: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    prenorm: bool = False,
+    residual_in_fp32: bool = False,
     is_rms_norm: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    return LayerNormLinearQuantFunction.apply(
         x,
         norm_weight,
         norm_bias,
 def rms_norm_linear_quant(
     x: torch.Tensor,
+    norm_weight: torch.Tensor | None,
+    norm_bias: torch.Tensor | None,
     linear_weight: torch.Tensor,
+    linear_bias: torch.Tensor | None,
+    residual: torch.Tensor | None = None,
     eps: float = 1e-5,
     prenorm: bool = False,
     residual_in_fp32: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    return layer_norm_linear_quant(
         x=x,
         norm_weight=norm_weight,
         norm_bias=norm_bias,
     )
+def bit_linear(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_bias: torch.Tensor | None = None,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    """Apply RMS normalization and a linear projection with quantized activations and weights."""
+    return layer_norm_linear_quant(
+        x=x,
+        norm_weight=norm_weight,
+        norm_bias=norm_bias,
+        linear_weight=weight,
+        linear_bias=bias,
+        eps=eps,
         is_rms_norm=True,
     )
 class BitLinear(nn.Linear):
     """
+    RMS-normalized linear layer with 8-bit activations and ternary weights.
+    Quantization uses a straight-through estimator during training.
+    Efficient deployment requires specialized kernels.
+    Args:
+        in_features (int):
+            Size of each input sample.
+        out_features (int):
+            Size of each output sample.
+        bias (bool, Optional):
+            Whether to allocate an additive bias. Default: `False`.
+        norm_eps (float, Optional):
+            Epsilon for RMS normalization. Default: 1e-8.
+    """
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, norm_eps: float = 1e-8) -> None:
+        super().__init__(in_features=in_features, out_features=out_features, bias=bias)
+        self.norm = RMSNorm(hidden_size=in_features, eps=norm_eps, dtype=torch.float32)
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({super().extra_repr()}, norm_eps={self.norm.eps})"
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         w = self.weight
+        x_norm = self.norm(x=x)
+        # straight-through gradients bypass quantization
+        x_quant = x_norm + (activation_quant(x=x_norm) - x_norm).detach()
+        w_quant = w + (weight_quant(w=w) - w).detach()
+        y = F.linear(input=x_quant, weight=w_quant)
         return y
 class FusedBitLinear(BitLinear):
+    """BitLinear with fused RMS normalization and activation quantization."""
+    def __init__(self, in_features: int, out_features: int, bias: bool = False) -> None:
+        super().__init__(in_features=in_features, out_features=out_features, bias=bias)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return layer_norm_linear_quant(
+            x=x,
+            norm_weight=self.norm.weight,
+            norm_bias=self.norm.bias,
+            linear_weight=self.weight,
+            linear_bias=self.bias,
+            eps=self.norm.eps,
             is_rms_norm=True,
         )

build/torch-cuda/modules/fused_cross_entropy.py DELETED Viewed

@@ -1,459 +0,0 @@
-# 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 typing import Any
-import torch
-import torch.nn as nn
-import triton
-import triton.language as tl
-from ..modules.backends import dispatch
-from ..ops.utils.op import exp, log, tanh
-from ..utils import input_guard
-# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for
-# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent
-# version of PyTorch. The following 2 lines are for backward compatibility with
-# older PyTorch.
-if "all_gather_into_tensor" not in dir(torch.distributed):
-    torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base
-@triton.heuristics({
-    "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0,
-    "HAS_SOFTCAPPING": lambda args: args["logit_softcapping"] is not None,
-})
-@triton.jit
-def cross_entropy_fwd_kernel(
-    loss_ptr,  # data ptrs
-    lse_ptr,
-    z_loss_ptr,
-    logits_ptr,
-    labels_ptr,
-    label_smoothing,
-    logit_scale,
-    lse_square_scale,
-    logit_softcapping,
-    ignore_index,
-    total_classes,
-    class_start_idx,  # Useful for tensor parallel when each rank only has a subset of classes
-    n_cols,  # shapes
-    n_rows,
-    logits_row_stride,  # strides
-    BLOCK_SIZE: tl.constexpr,
-    HAS_SMOOTHING: tl.constexpr,
-    HAS_SOFTCAPPING: tl.constexpr,
-    # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE
-    SPLIT: tl.constexpr,
-):
-    row_idx = tl.program_id(0)
-    col_block_idx = tl.program_id(1)
-    logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64)
-    col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    label_idx = tl.load(labels_ptr + row_idx)
-    logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf"))
-    logits = logits.to(tl.float32) * logit_scale
-    if HAS_SOFTCAPPING:
-        logits = logit_softcapping * tanh(logits / logit_softcapping)
-    max_logits = tl.max(logits, 0)
-    if HAS_SMOOTHING:
-        sum_logits = tl.sum(tl.where(col_offsets < n_cols, logits, 0.0), 0)
-    lse = log(tl.sum(exp(logits - max_logits), 0)) + max_logits
-    tl.store(lse_ptr + col_block_idx * n_rows + row_idx, lse)
-    if label_idx == ignore_index:
-        loss = 0.0
-        z_loss = 0.0
-    else:
-        label_idx -= class_start_idx
-        if label_idx >= col_block_idx * BLOCK_SIZE and label_idx < min(
-            n_cols, (col_block_idx + 1) * BLOCK_SIZE,
-        ):
-            logits_label = tl.load(logits_ptr + label_idx).to(tl.float32) * logit_scale
-            if HAS_SOFTCAPPING:
-                logits_label = logit_softcapping * tanh(logits_label / logit_softcapping)
-            if HAS_SMOOTHING:
-                loss = (
-                    (lse if not SPLIT else 0.0)
-                    - label_smoothing * sum_logits / total_classes
-                    - (1 - label_smoothing) * logits_label
-                )
-            else:
-                loss = (lse if not SPLIT else 0.0) - logits_label
-        else:
-            # If label is out of bounds, we set the CE loss to 0.0. But we still want the label_smoothing loss
-            if HAS_SMOOTHING:
-                loss = label_smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes)
-            else:
-                loss = 0.0
-        if not SPLIT:
-            z_loss = lse_square_scale * lse * lse
-            loss += z_loss
-        else:
-            z_loss = 0.0
-    tl.store(loss_ptr + col_block_idx * n_rows + row_idx, loss)
-    if not SPLIT:
-        tl.store(z_loss_ptr + col_block_idx * n_rows + row_idx, z_loss)
-@triton.heuristics({
-    "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0,
-    "HAS_SOFTCAPPING": lambda args: args["logit_softcapping"] is not None,
-})
-@triton.jit
-def cross_entropy_bwd_kernel(
-    dlogits_ptr,  # data ptrs
-    dloss_ptr,
-    logits_ptr,
-    lse_ptr,
-    labels_ptr,
-    label_smoothing,
-    logit_scale,
-    lse_square_scale,
-    logit_softcapping,
-    ignore_index,
-    total_classes,
-    class_start_idx,  # Useful for tensor parallel when each rank only has a subset of classes
-    n_cols,  # shapes
-    logits_row_stride,  # strides
-    dlogits_row_stride,
-    dloss_row_stride,
-    BLOCK_SIZE: tl.constexpr,
-    HAS_SMOOTHING: tl.constexpr,
-    HAS_SOFTCAPPING: tl.constexpr,
-):
-    row_idx = tl.program_id(0)
-    col_block_idx = tl.program_id(1)
-    logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64)
-    dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64)
-    col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    label_idx = tl.load(labels_ptr + row_idx)
-    if label_idx != ignore_index:
-        dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride)
-    else:
-        dloss = 0.0
-    logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to(
-        tl.float32,
-    ) * logit_scale
-    if HAS_SOFTCAPPING:
-        t = tanh(logits / logit_softcapping)
-        logits = logit_softcapping * t
-    lse = tl.load(lse_ptr + row_idx)
-    probs = exp(logits - lse)
-    probs += 2.0 * lse_square_scale * lse * probs
-    label_idx -= class_start_idx
-    if HAS_SMOOTHING:
-        smooth_negative = label_smoothing / total_classes
-        probs = tl.where(col_offsets == label_idx, probs - (1 - label_smoothing), probs) - smooth_negative
-    else:
-        probs = tl.where(col_offsets == label_idx, probs - 1.0, probs)
-    # d(softcap * tanh(x/softcap))/dx = 1 - tanh(x/softcap)^2
-    if HAS_SOFTCAPPING:
-        probs = probs * (1.0 - t * t)
-    tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols)
-def fused_cross_entropy_forward(
-    logits: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float = 0.0,
-    logit_scale: float = 1.0,
-    lse_square_scale: float = 0.0,
-    logit_softcapping: float = None,
-    ignore_index: int = -100,
-    process_group=None,
-):
-    n_rows, n_cols = logits.shape
-    assert target.shape == (n_rows,)
-    world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group)
-    total_classes = world_size * n_cols
-    rank = 0 if process_group is None else torch.distributed.get_rank(process_group)
-    class_start_idx = rank * n_cols
-    if logits.stride(-1) != 1:
-        logits = logits.contiguous()
-    # Set these similar to https://github.com/triton-lang/triton/blob/main/python/tutorials/02-fused-softmax.py
-    MAX_BLOCK_SIZE = 64 * 1024
-    BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE)
-    num_warps = (
-        4
-        if BLOCK_SIZE < 2048
-        else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32))
-    )
-    # We may split the lse computation across multiple blocks, then do a reduction
-    # lse(local_lse) to get the final LSE. This is faster for large n_cols (e.g., > 64k)
-    # where having just one thread block processing more than 64k elements is slow.
-    split = world_size > 1 or n_cols > MAX_BLOCK_SIZE
-    n_splits = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE
-    loss_shape = (n_splits, n_rows) if n_splits > 1 else (n_rows,)
-    losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device)
-    lse = torch.empty(*loss_shape, dtype=torch.float, device=logits.device)
-    z_losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device)
-    cross_entropy_fwd_kernel[(n_rows, n_splits)](
-        losses,  # data ptrs
-        lse,
-        z_losses,
-        logits,
-        target,
-        label_smoothing,
-        logit_scale,
-        lse_square_scale,
-        logit_softcapping,
-        ignore_index,
-        total_classes,
-        class_start_idx,
-        n_cols,  # shapes
-        n_rows,
-        logits.stride(0),  # strides
-        BLOCK_SIZE=BLOCK_SIZE,  # constants
-        num_warps=num_warps,
-        SPLIT=split,
-    )
-    if split:
-        # If there's no label_smoothing, if target are in the vocab of this partition, losses contains
-        # - predicted logit, and 0 otherwise.
-        # If there's label_smoothing=0.1, for target in the vocab of this partition, losses contains
-        # -0.9 * predicted logit - 0.1 * sum logit / total_classes.
-        # For target not in the vocab of this partition, losses contains
-        # -0.1 * sum logit / total_classes.
-        if n_splits > 1:
-            lse = torch.logsumexp(lse, dim=0)
-            losses = losses.sum(dim=0)
-        if world_size > 1:
-            lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device)
-            torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group)
-            handle_losses = torch.distributed.all_reduce(
-                losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True,
-            )
-            lse = torch.logsumexp(lse_allgather, dim=0)
-            handle_losses.wait()
-        # After the allreduce, if there's no label_smoothing, the total losses are - predicted_logit,
-        # we just have to add the (global) lse.
-        # If there's label_smoothing=0.1, the total losses are
-        # -0.9 * predicted_logit - 0.1 * sum logit / total_classes.
-        # Again, we just have to add the (global) lse.
-        losses += lse
-        if lse_square_scale != 0.0:
-            z_losses = lse_square_scale * lse.square()
-            z_losses.masked_fill_(target == ignore_index, 0.0)
-            losses += z_losses
-        else:
-            z_losses = torch.zeros_like(losses)
-        losses.masked_fill_(target == ignore_index, 0.0)
-    return losses, z_losses, lse, total_classes, class_start_idx
-class CrossEntropyLossFunction(torch.autograd.Function):
-    @staticmethod
-    @input_guard
-    def forward(
-        ctx,
-        logits,
-        target,
-        label_smoothing=0.0,
-        logit_scale=1.0,
-        lse_square_scale=0.0,
-        logit_softcapping=None,
-        ignore_index=-100,
-        inplace_backward=False,
-        process_group=None,
-    ):
-        losses, z_losses, lse, total_classes, class_start_idx = fused_cross_entropy_forward(
-            logits,
-            target,
-            label_smoothing,
-            logit_scale,
-            lse_square_scale,
-            logit_softcapping,
-            ignore_index,
-            process_group,
-        )
-        ctx.save_for_backward(logits, lse, target)
-        ctx.mark_non_differentiable(z_losses)
-        ctx.label_smoothing = label_smoothing
-        ctx.logit_scale = logit_scale
-        ctx.lse_square_scale = lse_square_scale
-        ctx.logit_softcapping = logit_softcapping
-        ctx.ignore_index = ignore_index
-        ctx.total_classes = total_classes
-        ctx.class_start_idx = class_start_idx
-        ctx.inplace_backward = inplace_backward
-        return losses, z_losses
-    @staticmethod
-    @input_guard
-    def backward(ctx, grad_losses, grad_z_losses):
-        del grad_z_losses  # z_losses are only for logging.
-        logits, lse, target = ctx.saved_tensors
-        dlogits = logits if ctx.inplace_backward else torch.empty_like(logits)
-        n_rows, n_cols = logits.shape
-        BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024)
-        num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16)
-        def grid(META): return (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"]))  # noqa
-        cross_entropy_bwd_kernel[grid](
-            dlogits,  # data ptrs
-            grad_losses,
-            logits,
-            lse,
-            target,
-            ctx.label_smoothing,
-            ctx.logit_scale,
-            ctx.lse_square_scale,
-            ctx.logit_softcapping,
-            ctx.ignore_index,
-            ctx.total_classes,
-            ctx.class_start_idx,
-            n_cols,  # shapes
-            logits.stride(0),  # strides
-            dlogits.stride(0),
-            grad_losses.stride(0),
-            BLOCK_SIZE=BLOCK_SIZE,  # constants
-            num_warps=num_warps,
-        )
-        return dlogits, None, None, None, None, None, None, None, None, None
-@dispatch('modules')
-def cross_entropy_loss(
-    logits: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float = 0.0,
-    logit_scale: float = 1.0,
-    lse_square_scale: float = 0.0,
-    logit_softcapping: float = None,
-    ignore_index=-100,
-    inplace_backward: bool = False,
-    process_group=None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Arguments:
-        logits: [batch, vocab_size]
-        target: [batch,]
-        label_smoothing: float
-        logit_scale: float.
-            Multiply logits by this scale before calculating the loss.
-        lse_square_scale: float.
-            If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss.
-            This is also referred to as "z-loss".
-        logit_softcapping: float.
-            If > 0, apply logit softcapping: logits = softcap * tanh(logits / softcap).
-            This prevents logit magnitudes from growing unboundedly.
-        ignore_index: int.
-            If target == ignore_index, the loss is set to 0.0.
-        inplace_backward: bool.
-            If True, we do the backward pass in-place by modifying the logits.
-            This saves memory.
-        process_group:
-            if not None, we're doing Tensor Parallel: each process is responsible for
-            one part of the vocab. The loss will be aggregated across processes.
-    Returns:
-        losses: [batch,], float
-        z_losses: [batch,], float
-    """
-    return CrossEntropyLossFunction.apply(
-        logits,
-        target,
-        label_smoothing,
-        logit_scale,
-        lse_square_scale,
-        logit_softcapping,
-        ignore_index,
-        inplace_backward,
-        process_group,
-    )
-class FusedCrossEntropyLoss(nn.Module):
-    def __init__(
-        self,
-        ignore_index: int = -100,
-        reduction: str = "mean",
-        label_smoothing: float = 0.0,
-        logit_scale: float = 1.0,
-        lse_square_scale: float = 0.0,
-        logit_softcapping: float = None,
-        inplace_backward: bool = False,
-        process_group: Any = None,
-        return_z_loss: bool = False,
-    ):
-        """
-        Arguments:
-            ignore_index: int. If target == ignore_index, the loss is set to 0.0.
-            label_smoothing: float
-            lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss.
-                This is also referred to as "z-loss".
-            logit_softcapping: float. If > 0, apply logit softcapping:
-                logits = softcap * tanh(logits / softcap).
-                This prevents logit magnitudes from growing unboundedly.
-            inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits.
-                This saves memory.
-            process_group: if not None, we're doing Tensor Parallel: each process is responsible for
-                one part of the vocab. The loss will be aggregated across processes.
-            return_z_loss: bool. If True, we return the component of the loss contributed by
-                the lse_square_scale value. This value is only for logging and does not support
-                backprop.
-        """
-        super().__init__()
-        if reduction not in ["mean", "none", "sum"]:
-            raise NotImplementedError("Only support reduction = 'mean' or 'none' or 'sum'")
-        self.ignore_index = ignore_index
-        self.reduction = reduction
-        self.label_smoothing = label_smoothing
-        self.logit_scale = logit_scale
-        self.lse_square_scale = lse_square_scale
-        self.logit_softcapping = logit_softcapping
-        self.inplace_backward = inplace_backward
-        self.process_group = process_group
-        self.return_z_loss = return_z_loss
-    def forward(self, input, target):
-        """
-        Arguments:
-            input: (batch, vocab_size)
-            target: (batch,)
-        Returns:
-            losses: (batch,) if reduction is 'none', else (1,), dtype float
-            z_loss: (batch,) if reduction is 'none', else (1,), dtype float (if self.return_z_loss)
-        """
-        assert input.device.type in ('cuda', 'npu') and target.device.type in ('cuda', 'npu'), (
-            "Only support CUDA/NPU tensors"
-        )
-        loss, z_loss = cross_entropy_loss(
-            input,
-            target,
-            label_smoothing=self.label_smoothing,
-            logit_scale=self.logit_scale,
-            lse_square_scale=self.lse_square_scale,
-            logit_softcapping=self.logit_softcapping,
-            ignore_index=self.ignore_index,
-            inplace_backward=self.inplace_backward,
-            process_group=self.process_group,
-        )
-        if self.reduction == "mean":
-            loss = loss.sum() / (target != self.ignore_index).sum()
-        elif self.reduction == "sum":
-            loss = loss.sum()
-        else:
-            loss = loss
-        if not self.return_z_loss:
-            return loss
-        if self.reduction == "mean":
-            z_loss = z_loss.sum() / (target != self.ignore_index).sum()
-        elif self.reduction == "sum":
-            z_loss = z_loss.sum()
-        else:
-            z_loss = z_loss
-        return loss, z_loss

build/torch-cuda/modules/fused_cross_entropy/__init__.py ADDED Viewed

	@@ -0,0 +1,15 @@

+# 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 ...backends import GluonBackend
+from ...modules.fused_cross_entropy.ops import FusedCrossEntropyLoss, cross_entropy_loss
+__all__ = ['FusedCrossEntropyLoss', 'cross_entropy_loss']
+if GluonBackend.is_available():
+    from ...modules.fused_cross_entropy import gluon  # noqa: F401

build/torch-cuda/modules/fused_cross_entropy/gluon.py ADDED Viewed

	@@ -0,0 +1,402 @@

+# 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 __future__ import annotations
+import os
+from typing import TYPE_CHECKING
+import torch
+import triton
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language import extra
+from ...backends import GluonBackend, register
+from ...modules.fused_cross_entropy.ops import cross_entropy_bwd, cross_entropy_fwd
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
+    exp = extra.libdevice.fast_expf
+    log = extra.libdevice.fast_logf
+else:
+    exp = gl.exp
+    log = gl.log
+tanh = extra.libdevice.tanh
+@gluon.jit(do_not_specialize=['N'])
+def cross_entropy_fwd_kernel_gluon(
+    logits,  # [N, V]
+    target,  # [N]
+    loss,  # [NV, N] or [N] when NV == 1
+    z_loss,  # [N]
+    lse,  # [NV, N] or [N] when NV == 1
+    s_logits,
+    label_smoothing: gl.constexpr,
+    logit_scale: gl.constexpr,
+    lse_square_scale: gl.constexpr,
+    logit_softcapping: gl.constexpr,
+    ignore_index: gl.constexpr,
+    N,
+    V: gl.constexpr,
+    V_TOTAL: gl.constexpr,
+    V_START: gl.constexpr,
+    BV: gl.constexpr,
+    SPLIT: gl.constexpr,
+):
+    NW: gl.constexpr = 4
+    R_MAX: gl.constexpr = 128 // logits.dtype.element_ty.primitive_bitwidth
+    R: gl.constexpr = (V & -V) if (V & -V) < R_MAX else R_MAX
+    layout: gl.constexpr = gl.BlockedLayout([R], [32], [NW], [0])
+    i_n = gl.program_id(0).to(gl.int64)
+    i_v = gl.program_id(1).to(gl.int64)
+    o_v = i_v * BV + gl.arange(0, BV, layout=layout).to(gl.int64)
+    m_v = o_v < V
+    p_logits = logits + i_n * s_logits
+    # [BV]
+    b_logits = gl.load(p_logits + o_v, mask=m_v, other=-float('inf')).to(gl.float32) * logit_scale
+    if logit_softcapping is not None:
+        b_logits = logit_softcapping * tanh(b_logits / logit_softcapping)
+    # these transforms do not preserve -inf in padded lanes
+    if logit_softcapping is not None or logit_scale <= 0:
+        b_logits = gl.where(m_v, b_logits, -float('inf'))
+    if label_smoothing > 0:
+        b_sum = gl.sum(gl.where(m_v, b_logits, 0.0), 0)
+    b_max = gl.max(b_logits, 0)
+    b_lse = log(gl.sum(exp(b_logits - b_max), 0)) + b_max
+    b_target = gl.load(target + i_n).to(gl.int64)
+    b_loss = 0.0
+    b_z_loss = 0.0
+    if b_target != ignore_index:
+        b_target -= V_START
+        m_target = (b_target >= i_v * BV) & (b_target < gl.minimum(V, (i_v + 1) * BV))
+        b_target_logit = gl.load(p_logits + b_target, mask=m_target, other=0).to(gl.float32) * logit_scale
+        if logit_softcapping is not None:
+            b_target_logit = logit_softcapping * tanh(b_target_logit / logit_softcapping)
+        b_loss = -b_target_logit
+        if label_smoothing > 0:
+            b_loss = (1 - label_smoothing) * b_loss - label_smoothing * b_sum / V_TOTAL
+        if not SPLIT:
+            b_z_loss = lse_square_scale * b_lse * b_lse
+            b_loss += b_lse + b_z_loss
+    gl.store(loss + i_v * N + i_n, b_loss)
+    gl.store(lse + i_v * N + i_n, b_lse)
+    if not SPLIT:
+        gl.store(z_loss + i_n, b_z_loss)
+@gluon.jit
+def cross_entropy_bwd_kernel_gluon(
+    logits,  # [N, V]
+    target,  # [N]
+    lse,  # [N]
+    dloss,  # [N]
+    dlogits,  # [N, V]
+    s_logits,
+    s_dloss,
+    s_dlogits,
+    label_smoothing: gl.constexpr,
+    logit_scale: gl.constexpr,
+    lse_square_scale: gl.constexpr,
+    logit_softcapping: gl.constexpr,
+    ignore_index: gl.constexpr,
+    V: gl.constexpr,
+    V_TOTAL: gl.constexpr,
+    V_START: gl.constexpr,
+    BV: gl.constexpr,
+):
+    NW: gl.constexpr = 4
+    R_MAX: gl.constexpr = 128 // logits.dtype.element_ty.primitive_bitwidth
+    R: gl.constexpr = (V & -V) if (V & -V) < R_MAX else R_MAX
+    layout: gl.constexpr = gl.BlockedLayout([R], [32], [NW], [0])
+    i_n = gl.program_id(0).to(gl.int64)
+    i_v = gl.program_id(1).to(gl.int64)
+    o_local = gl.arange(0, BV, layout=layout)
+    o_v = i_v * BV + o_local.to(gl.int64)
+    if R == R_MAX and V > BV:
+        m_v = o_local < gl.minimum(V - i_v * BV, BV).to(gl.int32)
+    else:
+        m_v = o_v < V
+    p_logits = logits + i_n * s_logits + o_v
+    p_dlogits = dlogits + i_n * s_dlogits + o_v
+    b_target = gl.load(target + i_n).to(gl.int64)
+    b_dloss = gl.load(dloss + i_n * s_dloss, mask=b_target != ignore_index, other=0).to(gl.float32)
+    # [BV]
+    b_logits = gl.load(p_logits, mask=m_v, other=0).to(gl.float32) * logit_scale
+    if logit_softcapping is not None:
+        b_tanh = tanh(b_logits / logit_softcapping)
+        b_logits = logit_softcapping * b_tanh
+    b_lse = gl.load(lse + i_n)
+    # [BV]
+    b_probs = exp(b_logits - b_lse)
+    b_dlogits = b_probs + 2.0 * lse_square_scale * b_lse * b_probs
+    TARGET_WEIGHT: gl.constexpr = 1 - label_smoothing if label_smoothing > 0 else 1.0
+    b_target -= V_START
+    if R == R_MAX and V > BV:
+        b_target -= i_v * BV
+        if (b_target >= 0) & (b_target < BV):
+            b_dlogits -= gl.where(o_local == b_target.to(gl.int32), TARGET_WEIGHT, 0.0)
+    else:
+        b_dlogits -= gl.where(o_v == b_target, TARGET_WEIGHT, 0.0)
+    if label_smoothing > 0:
+        b_dlogits -= label_smoothing / V_TOTAL
+    if logit_softcapping is not None:
+        b_dlogits *= 1.0 - b_tanh * b_tanh
+    b_dlogits *= b_dloss * logit_scale
+    gl.store(p_dlogits, b_dlogits, mask=m_v)
+@gluon.jit(do_not_specialize=['N'])
+def cross_entropy_fwd_kernel_reduce_gluon(
+    partial_loss,
+    partial_lse,
+    target,
+    loss,
+    z_loss,
+    lse,
+    lse_square_scale: gl.constexpr,
+    ignore_index: gl.constexpr,
+    N,
+    NV: gl.constexpr,
+    BN: gl.constexpr,
+    BS: gl.constexpr,
+    SPLIT: gl.constexpr,
+    REDUCE_LOSS: gl.constexpr,
+):
+    # keep partition reductions within a warp to avoid shared memory
+    P: gl.constexpr = 128 // BN
+    layout: gl.constexpr = gl.BlockedLayout([1, 1], [P, 32 // P], [1, 4], [1, 0])
+    o_n = gl.program_id(0).to(gl.int64) * BN + gl.arange(0, BN, layout=gl.SliceLayout(0, layout))
+    o_s = gl.arange(0, BS, layout=gl.SliceLayout(1, layout)).to(gl.int64)
+    offsets = o_s[:, None] * N + o_n[None, :]
+    mask = (o_s[:, None] < NV) & (o_n[None, :] < N)
+    b_lse = gl.load(partial_lse + offsets, mask=mask, other=-float('inf')).to(gl.float32)
+    b_max = gl.max(b_lse, 0)
+    b_lse = b_max + log(gl.sum(exp(b_lse - b_max[None, :]), 0))
+    if REDUCE_LOSS:
+        b_loss = gl.sum(gl.load(partial_loss + offsets, mask=mask, other=0).to(gl.float32), 0)
+    else:
+        b_loss = gl.load(partial_loss + o_n, mask=o_n < N, other=0).to(gl.float32)
+    if not SPLIT:
+        valid = gl.load(target + o_n, mask=o_n < N, other=ignore_index) != ignore_index
+        b_z_loss = gl.where(valid, lse_square_scale * b_lse * b_lse, 0.)
+        b_loss = gl.where(valid, b_loss + b_lse + b_z_loss, 0.)
+        gl.store(z_loss + o_n, b_z_loss, mask=o_n < N)
+    gl.store(loss + o_n, b_loss, mask=o_n < N)
+    gl.store(lse + o_n, b_lse, mask=o_n < N)
+def cross_entropy_reduce_gluon(
+    loss: torch.Tensor,
+    z_loss: torch.Tensor,
+    lse: torch.Tensor,
+    target: torch.Tensor,
+    lse_square_scale: float,
+    ignore_index: int,
+    process_group: ProcessGroup | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    N = target.shape[0]
+    world_size = 1 if process_group is None else torch.distributed.get_world_size(group=process_group)
+    if lse.ndim == 2:
+        NV = lse.shape[0]
+        BS = triton.next_power_of_2(NV)
+        BN = 128 // max(1, min(8, BS // 4, 16384 // triton.next_power_of_2(N)))
+        reduced_loss = loss.new_empty(N)
+        reduced_lse = torch.empty_like(reduced_loss)
+        cross_entropy_fwd_kernel_reduce_gluon[(triton.cdiv(N, BN),)](
+            partial_loss=loss,
+            partial_lse=lse,
+            target=target,
+            loss=reduced_loss,
+            z_loss=z_loss,
+            lse=reduced_lse,
+            lse_square_scale=lse_square_scale,
+            ignore_index=ignore_index,
+            N=N,
+            NV=NV,
+            BN=BN,
+            BS=BS,
+            SPLIT=world_size > 1,
+            REDUCE_LOSS=True,
+            num_warps=4,
+        )
+        loss, lse = reduced_loss, reduced_lse
+    if world_size > 1:
+        BS = triton.next_power_of_2(world_size)
+        BN = 128 // max(1, min(8, BS // 4, 16384 // triton.next_power_of_2(N)))
+        gathered_lse = lse.new_empty(world_size, N)
+        torch.distributed.all_gather_into_tensor(output_tensor=gathered_lse, input_tensor=lse, group=process_group)
+        torch.distributed.all_reduce(tensor=loss, group=process_group)
+        cross_entropy_fwd_kernel_reduce_gluon[(triton.cdiv(N, BN),)](
+            partial_loss=loss,
+            partial_lse=gathered_lse,
+            target=target,
+            loss=loss,
+            z_loss=z_loss,
+            lse=lse,
+            lse_square_scale=lse_square_scale,
+            ignore_index=ignore_index,
+            N=N,
+            NV=world_size,
+            BN=BN,
+            BS=BS,
+            SPLIT=False,
+            REDUCE_LOSS=False,
+            num_warps=4,
+        )
+    return loss, z_loss, lse
+def cross_entropy_fwd_verifier(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    process_group: ProcessGroup | None = None,
+) -> tuple[bool, str | None]:
+    if logits.device.type != 'cuda' or target.device != logits.device:
+        return False, "Logits and targets must be on the same CUDA device"
+    if logits.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return False, "Gluon cross entropy supports fp16, bf16, and fp32 logits"
+    return True, None
+@register(cross_entropy_fwd, backend=GluonBackend, verifier=cross_entropy_fwd_verifier, env_var='FLA_CE_GLUON')
+def cross_entropy_fwd_gluon(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    process_group: ProcessGroup | None = None,
+):
+    N, V = logits.shape
+    assert target.shape == (N,)
+    world_size = 1 if process_group is None else torch.distributed.get_world_size(group=process_group)
+    total_classes = world_size * V
+    rank = 0 if process_group is None else torch.distributed.get_rank(group=process_group)
+    class_start_idx = rank * V
+    if logits.stride(-1) != 1:
+        logits = logits.contiguous()
+    BV = triton.next_power_of_2(V)
+    if logit_softcapping is not None:
+        BV = min(BV, 2048)
+    elif V > 16384:
+        BV = 8192
+    # vocab partitions contribute partial loss until the global LSE is available
+    NV = triton.cdiv(V, BV)
+    split = world_size > 1 or NV > 1
+    shape = (NV, N) if NV > 1 else (N,)
+    loss = logits.new_empty(shape, dtype=torch.float)
+    lse = torch.empty_like(loss)
+    z_loss = logits.new_empty(N, dtype=torch.float)
+    cross_entropy_fwd_kernel_gluon[(N, NV)](
+        logits=logits,
+        target=target,
+        loss=loss,
+        z_loss=z_loss,
+        lse=lse,
+        s_logits=logits.stride(0),
+        label_smoothing=label_smoothing,
+        logit_scale=logit_scale,
+        lse_square_scale=lse_square_scale,
+        logit_softcapping=logit_softcapping,
+        ignore_index=ignore_index,
+        N=N,
+        V=V,
+        V_TOTAL=total_classes,
+        V_START=class_start_idx,
+        BV=BV,
+        SPLIT=split,
+        num_warps=4,
+    )
+    if split:
+        loss, z_loss, lse = cross_entropy_reduce_gluon(
+            loss=loss,
+            z_loss=z_loss,
+            lse=lse,
+            target=target,
+            lse_square_scale=lse_square_scale,
+            ignore_index=ignore_index,
+            process_group=process_group,
+        )
+    return loss, z_loss, lse, total_classes, class_start_idx
+def cross_entropy_bwd_verifier(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    lse: torch.Tensor,
+    dloss: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    total_classes: int | None = None,
+    class_start_idx: int = 0,
+    inplace_backward: bool = False,
+) -> tuple[bool, str | None]:
+    return cross_entropy_fwd_verifier(logits=logits, target=target)
+@register(cross_entropy_bwd, backend=GluonBackend, verifier=cross_entropy_bwd_verifier, env_var='FLA_CE_GLUON')
+def cross_entropy_bwd_gluon(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    lse: torch.Tensor,
+    dloss: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    total_classes: int | None = None,
+    class_start_idx: int = 0,
+    inplace_backward: bool = False,
+) -> torch.Tensor:
+    N, V = logits.shape
+    BV = min(triton.next_power_of_2(V), 1024 if logit_softcapping is not None else 4096)
+    dlogits = logits if inplace_backward else torch.empty_like(logits)
+    cross_entropy_bwd_kernel_gluon[(N, triton.cdiv(V, BV))](
+        logits=logits,
+        target=target,
+        lse=lse,
+        dloss=dloss,
+        dlogits=dlogits,
+        s_logits=logits.stride(0),
+        s_dloss=dloss.stride(0),
+        s_dlogits=dlogits.stride(0),
+        label_smoothing=label_smoothing,
+        logit_scale=logit_scale,
+        lse_square_scale=lse_square_scale,
+        logit_softcapping=logit_softcapping,
+        ignore_index=ignore_index,
+        V=V,
+        V_TOTAL=V if total_classes is None else total_classes,
+        V_START=class_start_idx,
+        BV=BV,
+        num_warps=4,
+    )
+    return dlogits

build/torch-cuda/modules/fused_cross_entropy/ops.py ADDED Viewed

	@@ -0,0 +1,538 @@

+# 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 __future__ import annotations
+from typing import TYPE_CHECKING
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+from triton.language.extra.libdevice import tanh
+from ...backends import dispatch
+from ...ops.utils.op import exp, log
+from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+@triton.jit(do_not_specialize=['N'])
+def cross_entropy_fwd_kernel(
+    logits,  # [N, V]
+    target,  # [N]
+    loss,  # [NV, N] or [N] when NV == 1
+    z_loss,  # [N]
+    lse,  # [NV, N] or [N] when NV == 1
+    s_logits,
+    label_smoothing: tl.constexpr,
+    logit_scale: tl.constexpr,
+    lse_square_scale: tl.constexpr,
+    logit_softcapping: tl.constexpr,
+    ignore_index: tl.constexpr,
+    N,
+    V: tl.constexpr,
+    V_TOTAL: tl.constexpr,
+    V_START: tl.constexpr,
+    BV: tl.constexpr,
+    SPLIT: tl.constexpr,
+):
+    i_n = tl.program_id(0).to(tl.int64)
+    i_v = tl.program_id(1)
+    o_v = i_v * BV + tl.arange(0, BV)
+    m_v = o_v < V
+    p_logits = logits + i_n * s_logits
+    # [BV]
+    b_logits = tl.load(p_logits + o_v, mask=m_v, other=-float('inf')).to(tl.float32) * logit_scale
+    if logit_softcapping is not None:
+        b_logits = logit_softcapping * tanh(b_logits / logit_softcapping)
+    # these transforms do not preserve -inf in padded lanes
+    if logit_softcapping is not None or logit_scale <= 0:
+        b_logits = tl.where(m_v, b_logits, -float('inf'))
+    if label_smoothing > 0:
+        b_sum = tl.sum(tl.where(m_v, b_logits, 0.0), 0)
+    b_max = tl.max(b_logits, 0)
+    b_lse = log(tl.sum(exp(b_logits - b_max), 0)) + b_max
+    b_target = tl.load(target + i_n).to(tl.int64)
+    b_loss = 0.0
+    b_z_loss = 0.0
+    if b_target != ignore_index:
+        b_target -= V_START
+        m_target = (b_target >= i_v * BV) & (b_target < tl.minimum(V, (i_v + 1) * BV))
+        b_target_logit = tl.load(p_logits + b_target, mask=m_target, other=0).to(tl.float32) * logit_scale
+        if logit_softcapping is not None:
+            b_target_logit = logit_softcapping * tanh(b_target_logit / logit_softcapping)
+        b_loss = -b_target_logit
+        if label_smoothing > 0:
+            b_loss = (1 - label_smoothing) * b_loss - label_smoothing * b_sum / V_TOTAL
+        if not SPLIT:
+            b_z_loss = lse_square_scale * b_lse * b_lse
+            b_loss += b_lse + b_z_loss
+    tl.store(loss + i_v * N + i_n, b_loss)
+    tl.store(lse + i_v * N + i_n, b_lse)
+    if not SPLIT:
+        tl.store(z_loss + i_n, b_z_loss)
+@triton.jit
+def cross_entropy_bwd_kernel(
+    logits,  # [N, V]
+    target,  # [N]
+    lse,  # [N]
+    dloss,  # [N]
+    dlogits,  # [N, V]
+    s_logits,
+    s_dloss,
+    s_dlogits,
+    label_smoothing: tl.constexpr,
+    logit_scale: tl.constexpr,
+    lse_square_scale: tl.constexpr,
+    logit_softcapping: tl.constexpr,
+    ignore_index: tl.constexpr,
+    V: tl.constexpr,
+    V_TOTAL: tl.constexpr,
+    V_START: tl.constexpr,
+    BV: tl.constexpr,
+):
+    i_n = tl.program_id(0).to(tl.int64)
+    i_v = tl.program_id(1).to(tl.int64)
+    o_v = i_v * BV + tl.arange(0, BV).to(tl.int64)
+    m_v = o_v < V
+    p_logits = logits + i_n * s_logits + o_v
+    p_dlogits = dlogits + i_n * s_dlogits + o_v
+    b_target = tl.load(target + i_n).to(tl.int64)
+    b_dloss = tl.load(dloss + i_n * s_dloss, mask=b_target != ignore_index, other=0).to(tl.float32)
+    # [BV]
+    b_logits = tl.load(p_logits, mask=m_v, other=0).to(tl.float32) * logit_scale
+    if logit_softcapping is not None:
+        b_tanh = tanh(b_logits / logit_softcapping)
+        b_logits = logit_softcapping * b_tanh
+    b_lse = tl.load(lse + i_n)
+    # [BV]
+    b_probs = exp(b_logits - b_lse)
+    b_dlogits = b_probs + 2.0 * lse_square_scale * b_lse * b_probs
+    b_target -= V_START
+    if label_smoothing > 0:
+        b_dlogits -= tl.where(o_v == b_target, 1 - label_smoothing, 0.0)
+        b_dlogits -= label_smoothing / V_TOTAL
+    else:
+        b_dlogits -= tl.where(o_v == b_target, 1.0, 0.0)
+    if logit_softcapping is not None:
+        b_dlogits *= 1.0 - b_tanh * b_tanh
+    b_dlogits *= b_dloss * logit_scale
+    tl.store(p_dlogits, b_dlogits, mask=m_v)
+@triton.jit(do_not_specialize=['N'])
+def cross_entropy_fwd_kernel_reduce(
+    partial_loss,
+    partial_lse,
+    target,
+    loss,
+    z_loss,
+    lse,
+    lse_square_scale: tl.constexpr,
+    ignore_index: tl.constexpr,
+    N,
+    NV: tl.constexpr,
+    BN: tl.constexpr,
+    BS: tl.constexpr,
+    SPLIT: tl.constexpr,
+    REDUCE_LOSS: tl.constexpr,
+):
+    o_n = tl.program_id(0).to(tl.int64) * BN + tl.arange(0, BN)
+    o_s = tl.arange(0, BS).to(tl.int64)
+    offsets = o_s[:, None] * N + o_n[None, :]
+    mask = (o_s[:, None] < NV) & (o_n[None, :] < N)
+    b_lse = tl.load(partial_lse + offsets, mask=mask, other=-float('inf')).to(tl.float32)
+    b_max = tl.max(b_lse, 0)
+    b_lse = b_max + log(tl.sum(exp(b_lse - b_max[None, :]), 0))
+    if REDUCE_LOSS:
+        b_loss = tl.sum(tl.load(partial_loss + offsets, mask=mask, other=0).to(tl.float32), 0)
+    else:
+        b_loss = tl.load(partial_loss + o_n, mask=o_n < N, other=0).to(tl.float32)
+    if not SPLIT:
+        valid = tl.load(target + o_n, mask=o_n < N, other=ignore_index) != ignore_index
+        b_z_loss = tl.where(valid, lse_square_scale * b_lse * b_lse, 0.)
+        b_loss = tl.where(valid, b_loss + b_lse + b_z_loss, 0.)
+        tl.store(z_loss + o_n, b_z_loss, mask=o_n < N)
+    tl.store(loss + o_n, b_loss, mask=o_n < N)
+    tl.store(lse + o_n, b_lse, mask=o_n < N)
+def cross_entropy_reduce(
+    loss: torch.Tensor,
+    z_loss: torch.Tensor,
+    lse: torch.Tensor,
+    target: torch.Tensor,
+    lse_square_scale: float,
+    ignore_index: int,
+    process_group: ProcessGroup | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    N = target.shape[0]
+    world_size = 1 if process_group is None else torch.distributed.get_world_size(group=process_group)
+    if lse.ndim == 2:
+        NV = lse.shape[0]
+        reduced_loss = loss.new_empty(N)
+        reduced_lse = torch.empty_like(reduced_loss)
+        cross_entropy_fwd_kernel_reduce[(triton.cdiv(N, 128),)](
+            partial_loss=loss,
+            partial_lse=lse,
+            target=target,
+            loss=reduced_loss,
+            z_loss=z_loss,
+            lse=reduced_lse,
+            lse_square_scale=lse_square_scale,
+            ignore_index=ignore_index,
+            N=N,
+            NV=NV,
+            BN=128,
+            BS=triton.next_power_of_2(NV),
+            SPLIT=world_size > 1,
+            REDUCE_LOSS=True,
+            num_warps=4,
+        )
+        loss, lse = reduced_loss, reduced_lse
+    if world_size > 1:
+        gathered_lse = lse.new_empty(world_size, N)
+        torch.distributed.all_gather_into_tensor(output_tensor=gathered_lse, input_tensor=lse, group=process_group)
+        torch.distributed.all_reduce(tensor=loss, group=process_group)
+        cross_entropy_fwd_kernel_reduce[(triton.cdiv(N, 128),)](
+            partial_loss=loss,
+            partial_lse=gathered_lse,
+            target=target,
+            loss=loss,
+            z_loss=z_loss,
+            lse=lse,
+            lse_square_scale=lse_square_scale,
+            ignore_index=ignore_index,
+            N=N,
+            NV=world_size,
+            BN=128,
+            BS=triton.next_power_of_2(world_size),
+            SPLIT=False,
+            REDUCE_LOSS=False,
+            num_warps=4,
+        )
+    return loss, z_loss, lse
+@dispatch
+def cross_entropy_fwd(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    process_group: ProcessGroup | None = None,
+):
+    N, V = logits.shape
+    assert target.shape == (N,)
+    world_size = 1 if process_group is None else torch.distributed.get_world_size(group=process_group)
+    total_classes = world_size * V
+    rank = 0 if process_group is None else torch.distributed.get_rank(group=process_group)
+    class_start_idx = rank * V
+    if logits.stride(-1) != 1:
+        logits = logits.contiguous()
+    BV = min(triton.next_power_of_2(V), 64 * 1024)
+    # reduce softcap register pressure without introducing an additional split reduction
+    if logit_softcapping is not None and V > BV:
+        BV = 8 * 1024
+    num_warps = 4 if BV < 2048 else (8 if BV < 8192 else 16)
+    # vocab partitions contribute partial loss until the global LSE is available
+    NV = triton.cdiv(V, BV)
+    split = world_size > 1 or NV > 1
+    shape = (NV, N) if NV > 1 else (N,)
+    loss = logits.new_empty(shape, dtype=torch.float)
+    lse = torch.empty_like(loss)
+    z_loss = logits.new_empty(N, dtype=torch.float)
+    cross_entropy_fwd_kernel[(N, NV)](
+        logits=logits,
+        target=target,
+        loss=loss,
+        z_loss=z_loss,
+        lse=lse,
+        s_logits=logits.stride(0),
+        label_smoothing=label_smoothing,
+        logit_scale=logit_scale,
+        lse_square_scale=lse_square_scale,
+        logit_softcapping=logit_softcapping,
+        ignore_index=ignore_index,
+        N=N,
+        V=V,
+        V_TOTAL=total_classes,
+        V_START=class_start_idx,
+        BV=BV,
+        SPLIT=split,
+        num_warps=num_warps,
+    )
+    if split:
+        loss, z_loss, lse = cross_entropy_reduce(
+            loss=loss,
+            z_loss=z_loss,
+            lse=lse,
+            target=target,
+            lse_square_scale=lse_square_scale,
+            ignore_index=ignore_index,
+            process_group=process_group,
+        )
+    return loss, z_loss, lse, total_classes, class_start_idx
+@dispatch
+def cross_entropy_bwd(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    lse: torch.Tensor,
+    dloss: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    total_classes: int | None = None,
+    class_start_idx: int = 0,
+    inplace_backward: bool = False,
+):
+    N, V = logits.shape
+    BV = min(triton.next_power_of_2(V), 4 * 1024)
+    dlogits = logits if inplace_backward else torch.empty_like(logits)
+    cross_entropy_bwd_kernel[(N, triton.cdiv(V, BV))](
+        logits=logits,
+        target=target,
+        lse=lse,
+        dloss=dloss,
+        dlogits=dlogits,
+        s_logits=logits.stride(0),
+        s_dloss=dloss.stride(0),
+        s_dlogits=dlogits.stride(0),
+        label_smoothing=label_smoothing,
+        logit_scale=logit_scale,
+        lse_square_scale=lse_square_scale,
+        logit_softcapping=logit_softcapping,
+        ignore_index=ignore_index,
+        V=V,
+        V_TOTAL=V if total_classes is None else total_classes,
+        V_START=class_start_idx,
+        BV=BV,
+        num_warps=4 if BV < 2048 else 8,
+    )
+    return dlogits
+class FusedCrossEntropyFunction(torch.autograd.Function):
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(
+        ctx,
+        logits,
+        target,
+        label_smoothing=0.0,
+        logit_scale=1.0,
+        lse_square_scale=0.0,
+        logit_softcapping=None,
+        ignore_index=-100,
+        inplace_backward=False,
+        process_group: ProcessGroup | None = None,
+    ):
+        loss, z_loss, lse, total_classes, class_start_idx = cross_entropy_fwd(
+            logits=logits,
+            target=target,
+            label_smoothing=label_smoothing,
+            logit_scale=logit_scale,
+            lse_square_scale=lse_square_scale,
+            logit_softcapping=logit_softcapping,
+            ignore_index=ignore_index,
+            process_group=process_group,
+        )
+        ctx.save_for_backward(logits, target, lse)
+        ctx.mark_non_differentiable(z_loss)
+        ctx.label_smoothing = label_smoothing
+        ctx.logit_scale = logit_scale
+        ctx.lse_square_scale = lse_square_scale
+        ctx.logit_softcapping = logit_softcapping
+        ctx.ignore_index = ignore_index
+        ctx.total_classes = total_classes
+        ctx.class_start_idx = class_start_idx
+        ctx.inplace_backward = inplace_backward
+        return loss, z_loss
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(ctx, dloss, dz_loss):
+        # z-loss is returned for logging and is already included in the differentiable loss
+        del dz_loss
+        logits, target, lse = ctx.saved_tensors
+        dlogits = cross_entropy_bwd(
+            logits=logits,
+            target=target,
+            lse=lse,
+            dloss=dloss,
+            label_smoothing=ctx.label_smoothing,
+            logit_scale=ctx.logit_scale,
+            lse_square_scale=ctx.lse_square_scale,
+            logit_softcapping=ctx.logit_softcapping,
+            ignore_index=ctx.ignore_index,
+            total_classes=ctx.total_classes,
+            class_start_idx=ctx.class_start_idx,
+            inplace_backward=ctx.inplace_backward,
+        )
+        return dlogits, None, None, None, None, None, None, None, None
+fused_cross_entropy_forward = cross_entropy_fwd
+CrossEntropyLossFunction = FusedCrossEntropyFunction
+@dispatch
+def cross_entropy_loss(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    logit_scale: float = 1.0,
+    lse_square_scale: float = 0.0,
+    logit_softcapping: float | None = None,
+    ignore_index: int = -100,
+    inplace_backward: bool = False,
+    process_group: ProcessGroup | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute unreduced cross entropy and its non-differentiable z-loss component.
+    Args:
+        logits (torch.Tensor):
+            Input logits of shape `[N, V]`.
+        target (torch.Tensor):
+            Global target indices of shape `[N]`.
+        label_smoothing (float, Optional):
+            Uniform label smoothing coefficient. Default: 0.0.
+        logit_scale (float, Optional):
+            Scale applied before softcapping. Default: 1.0.
+        lse_square_scale (float, Optional):
+            Coefficient of the squared logsumexp regularizer. Default: 0.0.
+        logit_softcapping (float, Optional):
+            Softcap applied as `softcap * tanh(logits / softcap)`. Default: `None`.
+        ignore_index (int, Optional):
+            Target index excluded from loss and gradients. Default: -100.
+        inplace_backward (bool, Optional):
+            Whether to overwrite logits with their gradients. Default: `False`.
+        process_group (ProcessGroup, Optional):
+            Group sharding the vocabulary into equal contiguous partitions. Default: `None`.
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]:
+            FP32 loss and z-loss tensors of shape `[N]`. The loss includes z-loss.
+    """
+    return FusedCrossEntropyFunction.apply(
+        logits,
+        target,
+        label_smoothing,
+        logit_scale,
+        lse_square_scale,
+        logit_softcapping,
+        ignore_index,
+        inplace_backward,
+        process_group,
+    )
+class FusedCrossEntropyLoss(nn.Module):
+    def __init__(
+        self,
+        ignore_index: int = -100,
+        reduction: str = "mean",
+        label_smoothing: float = 0.0,
+        logit_scale: float = 1.0,
+        lse_square_scale: float = 0.0,
+        logit_softcapping: float | None = None,
+        inplace_backward: bool = False,
+        process_group: ProcessGroup | None = None,
+        return_z_loss: bool = False,
+    ):
+        """Cross entropy with optional logit transforms and vocabulary parallelism.
+        Args:
+            ignore_index (int, Optional):
+                Target index excluded from loss and gradients. Default: -100.
+            reduction (str, Optional):
+                Reduction over non-ignored targets: `mean`, `sum`, or `none`. Default: `mean`.
+            label_smoothing (float, Optional):
+                Uniform label smoothing coefficient. Default: 0.0.
+            logit_scale (float, Optional):
+                Scale applied before softcapping. Default: 1.0.
+            lse_square_scale (float, Optional):
+                Coefficient of the squared logsumexp regularizer. Default: 0.0.
+            logit_softcapping (float, Optional):
+                Softcap applied as `softcap * tanh(logits / softcap)`. Default: `None`.
+            inplace_backward (bool, Optional):
+                Whether to overwrite logits with their gradients. Default: `False`.
+            process_group (ProcessGroup, Optional):
+                Group sharding the vocabulary into equal contiguous partitions. Default: `None`.
+            return_z_loss (bool, Optional):
+                Whether to also return z-loss for logging, without its own backward. Default: `False`.
+        """
+        super().__init__()
+        if reduction not in ('mean', 'sum', 'none'):
+            raise NotImplementedError(f"Unsupported reduction: {reduction}")
+        self.ignore_index = ignore_index
+        self.reduction = reduction
+        self.label_smoothing = label_smoothing
+        self.logit_scale = logit_scale
+        self.lse_square_scale = lse_square_scale
+        self.logit_softcapping = logit_softcapping
+        self.inplace_backward = inplace_backward
+        self.process_group = process_group
+        self.return_z_loss = return_z_loss
+    def forward(self, input: torch.Tensor, target: torch.Tensor):
+        """Compute loss from `[N, V]` logits and `[N]` targets.
+        Return FP32 scalars for `mean` or `sum`, and `[N]` tensors for `none`.
+        When `return_z_loss=True`, also return the reduced z-loss component.
+        """
+        assert input.device.type in ('cuda', 'npu', 'xpu') and target.device.type in ('cuda', 'npu', 'xpu'), (
+            "Only support CUDA/NPU/XPU tensors"
+        )
+        loss, z_loss = cross_entropy_loss(
+            logits=input,
+            target=target,
+            label_smoothing=self.label_smoothing,
+            logit_scale=self.logit_scale,
+            lse_square_scale=self.lse_square_scale,
+            logit_softcapping=self.logit_softcapping,
+            ignore_index=self.ignore_index,
+            inplace_backward=self.inplace_backward,
+            process_group=self.process_group,
+        )
+        if self.reduction != 'none':
+            loss = loss.sum()
+            if self.return_z_loss:
+                z_loss = z_loss.sum()
+            if self.reduction == 'mean':
+                total = (target != self.ignore_index).sum()
+                loss = loss / total
+                if self.return_z_loss:
+                    z_loss = z_loss / total
+        return (loss, z_loss) if self.return_z_loss else loss

build/torch-cuda/modules/fused_kl_div/__init__.py ADDED Viewed

	@@ -0,0 +1,11 @@

+# 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 ...modules.fused_kl_div.ops import FusedKLDivLoss, fused_kl_div_loss
+__all__ = ['FusedKLDivLoss', 'fused_kl_div_loss']

build/torch-cuda/modules/{fused_kl_div.py → fused_kl_div/ops.py} RENAMED Viewed

@@ -11,14 +11,14 @@ import torch.nn.functional as F
 import triton
 import triton.language as tl
-from ..modules.backends import dispatch
-from ..ops.utils.op import exp, log
-from ..utils import IS_AMD, input_guard
-# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576
 # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
-# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
-# The optimal maximum block size depends on your hardware, your kernel, and your dtype
 MAX_FUSED_SIZE = 65536 // 2
 STATIC_WARPS = 32 if not IS_AMD else 16
@@ -42,22 +42,19 @@ def kl_div_kernel(
     logits += i_n * s_logits
     target_logits += i_n * s_logits
-    # m is the max value. use the notation from the paper
     sm = float('-inf')
     tm = float('-inf')
-    # d is the sum. use the notation from the paper
     sd, td = 0.0, 0.0
     NV = tl.cdiv(V, BV)
     for iv in range(0, NV):
         o_x = iv * BV + tl.arange(0, BV)
-        # for student
         b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf'))
         b_sm = tl.max(b_sl)
         m_new = tl.maximum(sm, b_sm)
         sd = sd * exp(sm - m_new) + tl.sum(exp(b_sl - m_new))
         sm = m_new
-        # for teacher
         b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf'))
         b_tm = tl.max(b_tl)
         m_new = tl.maximum(tm, b_tm)
@@ -89,45 +86,28 @@ def kl_div_kernel(
 @triton.jit
-def elementwise_mul_kernel(
-    x,
-    g,
-    N: tl.constexpr,
-    B: tl.constexpr,
-):
-    """
-    This function multiplies each element of the tensor pointed by x with the value pointed by g.
-    The multiplication is performed in-place on the tensor pointed by x.
-    Parameters:
-    x:
-        Pointer to the input tensor.
-    g:
-        Pointer to the gradient output value.
-    N (int):
-        The number of columns in the input tensor.
-    B (int):
-        The block size for Triton operations.
-    """
-    # Get the program ID and convert it to int64 to avoid overflow
     i_x = tl.program_id(0).to(tl.int64)
     o_x = i_x * B + tl.arange(0, B)
-    # Load the gradient output value
     b_g = tl.load(g)
     b_x = tl.load(x + o_x, mask=o_x < N)
     tl.store(x + o_x, b_x * b_g, mask=o_x < N)
-@dispatch('modules')
-def fused_kl_div_forward(
     x: torch.Tensor,
     target_x: torch.Tensor,
     weight: torch.Tensor,
     target_weight: torch.Tensor,
     reduction: str = 'batchmean',
     accumulate_grad_in_fp32: bool = True,
 ):
     device = x.device
@@ -146,9 +126,8 @@ def fused_kl_div_forward(
     grad_dtype = torch.float32 if accumulate_grad_in_fp32 else weight.dtype
-    dx = torch.zeros_like(x, device=device)
-    dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None
-    # we use fp32 for loss accumulator
     loss = torch.zeros(N, dtype=torch.float32, device=device)
     for ic in range(NC):
@@ -156,14 +135,12 @@ def fused_kl_div_forward(
         # [C, N]
         c_sx = x[start:end]
         c_tx = target_x[start:end]
-        # when doing matmul, use the original precision
         # [C, V]
         c_sl = F.linear(c_sx, weight)
         c_tl = F.linear(c_tx, target_weight)
-        if weight is not None and c_sx.dtype != grad_dtype:
             c_sx = c_sx.to(dtype=grad_dtype)
-        # unreduced loss
         c_loss = loss[start:end]
         # Here we calculate the gradient of c_sx in place so we can save memory.
@@ -187,15 +164,11 @@ def fused_kl_div_forward(
         # Thus, we need an additional scaling factor of (n_non_ignore/total) to scale the gradients.
         # [C, H]
-        dx[start:end] = torch.mm(c_sl, weight)
-        if weight is not None:
-            torch.addmm(
-                input=dw,
-                mat1=c_sl.t().to(dtype=grad_dtype),
-                mat2=c_sx,
-                out=dw,
-            )
     loss = loss.sum()
     if dw is not None:
@@ -203,37 +176,15 @@ def fused_kl_div_forward(
     return loss, dx, dw
-@dispatch('modules')
-def fused_kl_div_backward(
-    do: torch.Tensor,
-    dx: torch.Tensor,
-    dw: torch.Tensor,
-):
-    # If cross entropy is the last layer, do is 1.0. Skip the mul to save time
-    if torch.ne(do, torch.tensor(1.0, device=do.device)):
-        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
-        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
-        N, H = dx.shape
         B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
-        elementwise_mul_kernel[(triton.cdiv(N * H, B),)](
-            x=dx,
-            g=do,
-            N=N*H,
-            B=B,
-            num_warps=STATIC_WARPS,
-        )
-        # handle dw
-        if dw is not None:
-            V, H = dw.shape
-            elementwise_mul_kernel[(triton.cdiv(V * H, B),)](
-                x=dw,
-                g=do,
-                N=V*H,
-                B=B,
-                num_warps=STATIC_WARPS,
-            )
     return dx, dw
@@ -250,14 +201,17 @@ class FusedKLDivLossFunction(torch.autograd.Function):
         target_weight: torch.Tensor,
         reduction: str,
         accumulate_grad_in_fp32: bool,
     ):
-        loss, dx, dw = fused_kl_div_forward(
             x=x,
             target_x=target_x,
             weight=weight,
             target_weight=target_weight,
             reduction=reduction,
             accumulate_grad_in_fp32=accumulate_grad_in_fp32,
         )
         ctx.save_for_backward(dx, dw)
         return loss
@@ -266,8 +220,8 @@ class FusedKLDivLossFunction(torch.autograd.Function):
     @input_guard
     def backward(ctx, do):
         dx, dw = ctx.saved_tensors
-        dx, dw = fused_kl_div_backward(do, dx, dw)
-        return dx, None, dw, None, None, None
 def fused_kl_div_loss(
@@ -280,23 +234,25 @@ def fused_kl_div_loss(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Args:
-        x (`torch.Tensor`):
             Tensor of shape `[batch_size * seq_len, hidden_size]`.
-        target_x (`torch.Tensor`):
             Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`.
             Must not require gradients.
-        weight (`torch.Tensor`):
             Tensor of shape `[vocab_size, hidden_size]`.
-        target_weight (`torch.Tensor`):
             Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`.
             Must not require gradients.
-        reduction (`str`):
             Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'.
-        accumulate_grad_in_fp32 (`bool`):
-            Whether to accumulate the student weight gradient in fp32 before casting it back
-            to `weight.dtype`. Default: `True`.
     Returns:
-        loss
     """
     if target_x.requires_grad or target_weight.requires_grad:
         raise RuntimeError(
@@ -311,26 +267,24 @@ def fused_kl_div_loss(
         target_weight,
         reduction,
         accumulate_grad_in_fp32,
     )
 class FusedKLDivLoss(nn.Module):
-    def __init__(
-        self,
-        reduction: str = 'batchmean',
-        accumulate_grad_in_fp32: bool = True,
-    ):
         """
         Args:
-            reduction (`str`):
                 Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'.
-            accumulate_grad_in_fp32 (`bool`):
-                Whether to accumulate the student weight gradient in fp32 before casting it back
-                to `weight.dtype`. Default: `True`.
         Note:
-            FusedKLDivLoss only computes gradients for `x` and `weight`; `target_x` and
-            `target_weight` are treated as frozen teacher tensors and must not require gradients.
         """
         super().__init__()
@@ -339,27 +293,23 @@ class FusedKLDivLoss(nn.Module):
         self.reduction = reduction
         self.accumulate_grad_in_fp32 = accumulate_grad_in_fp32
-    def forward(
-        self,
-        x: torch.Tensor,
-        target_x: torch.Tensor,
-        weight: torch.Tensor,
-        target_weight: torch.Tensor,
-    ):
         """
         Args:
-            x (`torch.Tensor`):
                 Tensor of shape `[batch_size * seq_len, hidden_size]`.
-            target_x (`torch.Tensor`):
                 Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`.
                 Must not require gradients.
-            weight (`torch.Tensor`):
                 Tensor of shape `[vocab_size, hidden_size]`.
-            target_weight (`torch.Tensor`):
                 Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`.
                 Must not require gradients.
         Returns:
-            loss
         """
         loss = fused_kl_div_loss(
             x=x,

 import triton
 import triton.language as tl
+from ...backends import dispatch
+from ...ops.utils.op import exp, log
+from ...utils import IS_AMD, input_guard
+# the hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576
 # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
+# however, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
+# the optimal maximum block size depends on your hardware, your kernel, and your dtype
 MAX_FUSED_SIZE = 65536 // 2
 STATIC_WARPS = 32 if not IS_AMD else 16
     logits += i_n * s_logits
     target_logits += i_n * s_logits
+    # running maxima and exponential sums for student and teacher logits.
     sm = float('-inf')
     tm = float('-inf')
     sd, td = 0.0, 0.0
     NV = tl.cdiv(V, BV)
     for iv in range(0, NV):
         o_x = iv * BV + tl.arange(0, BV)
         b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf'))
         b_sm = tl.max(b_sl)
         m_new = tl.maximum(sm, b_sm)
         sd = sd * exp(sm - m_new) + tl.sum(exp(b_sl - m_new))
         sm = m_new
         b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf'))
         b_tm = tl.max(b_tl)
         m_new = tl.maximum(tm, b_tm)
 @triton.jit
+def elementwise_mul_kernel(x, g, N: tl.constexpr, B: tl.constexpr):
     i_x = tl.program_id(0).to(tl.int64)
     o_x = i_x * B + tl.arange(0, B)
     b_g = tl.load(g)
+    if b_g == 1.0:
+        return
     b_x = tl.load(x + o_x, mask=o_x < N)
     tl.store(x + o_x, b_x * b_g, mask=o_x < N)
+@dispatch
+def fused_kl_div_fwd(
     x: torch.Tensor,
     target_x: torch.Tensor,
     weight: torch.Tensor,
     target_weight: torch.Tensor,
     reduction: str = 'batchmean',
     accumulate_grad_in_fp32: bool = True,
+    use_dx: bool = True,
+    use_dw: bool = True,
 ):
     device = x.device
     grad_dtype = torch.float32 if accumulate_grad_in_fp32 else weight.dtype
+    dx = torch.zeros_like(x, device=device) if use_dx else None
+    dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if use_dw else None
     loss = torch.zeros(N, dtype=torch.float32, device=device)
     for ic in range(NC):
         # [C, N]
         c_sx = x[start:end]
         c_tx = target_x[start:end]
         # [C, V]
         c_sl = F.linear(c_sx, weight)
         c_tl = F.linear(c_tx, target_weight)
+        if dw is not None and c_sx.dtype != grad_dtype:
             c_sx = c_sx.to(dtype=grad_dtype)
         c_loss = loss[start:end]
         # Here we calculate the gradient of c_sx in place so we can save memory.
         # Thus, we need an additional scaling factor of (n_non_ignore/total) to scale the gradients.
         # [C, H]
+        if dx is not None:
+            dx[start:end] = torch.mm(c_sl, weight)
+        if dw is not None:
+            torch.addmm(input=dw, mat1=c_sl.t().to(dtype=grad_dtype), mat2=c_sx, out=dw)
     loss = loss.sum()
     if dw is not None:
     return loss, dx, dw
+@dispatch
+def fused_kl_div_bwd(do: torch.Tensor, dx: torch.Tensor | None, dw: torch.Tensor | None):
+    # in-place gradient storage across repeated backward calls causes PyTorch anomalies; use Triton.
+    for grad in (dx, dw):
+        if grad is None:
+            continue
+        N, H = grad.shape
         B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
+        elementwise_mul_kernel[(triton.cdiv(N * H, B),)](x=grad, g=do, N=N*H, B=B, num_warps=STATIC_WARPS,)
     return dx, dw
         target_weight: torch.Tensor,
         reduction: str,
         accumulate_grad_in_fp32: bool,
+        grad_enabled: bool,
     ):
+        loss, dx, dw = fused_kl_div_fwd(
             x=x,
             target_x=target_x,
             weight=weight,
             target_weight=target_weight,
             reduction=reduction,
             accumulate_grad_in_fp32=accumulate_grad_in_fp32,
+            use_dx=grad_enabled and ctx.needs_input_grad[0],
+            use_dw=grad_enabled and ctx.needs_input_grad[2],
         )
         ctx.save_for_backward(dx, dw)
         return loss
     @input_guard
     def backward(ctx, do):
         dx, dw = ctx.saved_tensors
+        dx, dw = fused_kl_div_bwd(do=do, dx=dx, dw=dw)
+        return dx, None, dw, None, None, None, None
 def fused_kl_div_loss(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Args:
+        x (torch.Tensor):
             Tensor of shape `[batch_size * seq_len, hidden_size]`.
+        target_x (torch.Tensor):
             Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`.
             Must not require gradients.
+        weight (torch.Tensor):
             Tensor of shape `[vocab_size, hidden_size]`.
+        target_weight (torch.Tensor):
             Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`.
             Must not require gradients.
+        reduction (str, Optional):
             Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'.
+        accumulate_grad_in_fp32 (bool, Optional):
+            Whether to accumulate the student weight gradient in fp32 before casting it back to `weight.dtype`.
+            Default: `True`.
     Returns:
+        loss (torch.Tensor):
+            Scalar KL divergence loss.
     """
     if target_x.requires_grad or target_weight.requires_grad:
         raise RuntimeError(
         target_weight,
         reduction,
         accumulate_grad_in_fp32,
+        torch.is_grad_enabled(),
     )
 class FusedKLDivLoss(nn.Module):
+    def __init__(self, reduction: str = 'batchmean', accumulate_grad_in_fp32: bool = True):
         """
         Args:
+            reduction (str, Optional):
                 Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'.
+            accumulate_grad_in_fp32 (bool, Optional):
+                Whether to accumulate the student weight gradient in fp32 before casting it back to `weight.dtype`.
+                Default: `True`.
         Note:
+            FusedKLDivLoss only computes gradients for `x` and `weight`;
+            `target_x` and `target_weight` are treated as frozen teacher tensors and must not require gradients.
         """
         super().__init__()
         self.reduction = reduction
         self.accumulate_grad_in_fp32 = accumulate_grad_in_fp32
+    def forward(self, x: torch.Tensor, target_x: torch.Tensor, weight: torch.Tensor, target_weight: torch.Tensor):
         """
         Args:
+            x (torch.Tensor):
                 Tensor of shape `[batch_size * seq_len, hidden_size]`.
+            target_x (torch.Tensor):
                 Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`.
                 Must not require gradients.
+            weight (torch.Tensor):
                 Tensor of shape `[vocab_size, hidden_size]`.
+            target_weight (torch.Tensor):
                 Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`.
                 Must not require gradients.
         Returns:
+            loss (torch.Tensor):
+                Scalar KL divergence loss.
         """
         loss = fused_kl_div_loss(
             x=x,

build/torch-cuda/modules/fused_linear_cross_entropy/__init__.py ADDED Viewed

	@@ -0,0 +1,15 @@

+# 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 ...modules.fused_linear_cross_entropy.ops import (
+    FusedLinearCrossEntropyLoss,
+    LinearLossParallel,
+    fused_linear_cross_entropy_loss,
+)
+__all__ = ['FusedLinearCrossEntropyLoss', 'LinearLossParallel', 'fused_linear_cross_entropy_loss']

build/torch-cuda/modules/{fused_linear_cross_entropy.py → fused_linear_cross_entropy/ops.py} RENAMED Viewed

@@ -5,10 +5,12 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-# Code adapted from
-# https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py
 from functools import partial
 import torch
 import torch.nn as nn
@@ -19,19 +21,23 @@ from torch.distributed import DeviceMesh
 from torch.distributed.tensor import Replicate, Shard, distribute_module
 from torch.distributed.tensor.parallel import ParallelStyle
-from ..modules.backends import dispatch
-from ..ops.utils.op import exp, log, tanh
-from ..utils import IS_AMD, input_guard
 try:
     from torch.distributed.tensor import DTensor
 except (ImportError, AttributeError):
     DTensor = None
-# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576
 # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
-# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
-# The optimal maximum block size depends on your hardware, your kernel, and your dtype
 MAX_FUSED_SIZE = 65536 // 2
 STATIC_WARPS = 32 if not IS_AMD else 16
@@ -42,268 +48,86 @@ STATIC_WARPS = 32 if not IS_AMD else 16
 })
 @triton.jit
 def logsumexp_fwd_kernel(
-    x,
-    z,
     scale,
     softcapping,
     D: tl.constexpr,
-    B: tl.constexpr,
     HAS_SCALE: tl.constexpr,
     HAS_SOFTCAPPING: tl.constexpr,
 ):
     i_n, i_d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
-    o_d = i_d * B + tl.arange(0, B)
     m_d = o_d < D
     b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
     if HAS_SCALE:
         b_x = b_x * scale
     if HAS_SOFTCAPPING:
         b_x = softcapping * tanh(b_x / softcapping)
     b_m = tl.max(b_x, 0)
     b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m
-    tl.store(z + i_n * tl.cdiv(D, B) + i_d, b_z)
-@dispatch('modules')
-def logsumexp_fwd(
-    x,
-    scale: float | None = None,
-    softcapping: float | None = None,
-    dtype: torch.dtype | None = None,
 ):
     shape = x.shape
     x = x.view(-1, shape[-1])
     N, D = x.shape
-    B = min(triton.next_power_of_2(D), 64 * 1024)
-    ND = triton.cdiv(D, B)
     z = x.new_empty(N, ND, dtype=torch.float)
-    logsumexp_fwd_kernel[(N, ND)](
-        x=x,
-        z=z,
-        scale=scale,
-        softcapping=softcapping,
-        D=D,
-        B=B,
-    )
     z = z.logsumexp(-1).view(*shape[:-1])
     if dtype is not None and dtype != torch.float:
         z = z.to(dtype)
     return z
-@triton.jit
-def cross_entropy_kernel(
-    logits,
-    lse,
-    target,
-    loss,
-    total,
-    ignore_index,
-    label_smoothing: tl.constexpr,
-    logit_scale: tl.constexpr,
-    logit_softcapping: tl.constexpr,
-    reduction: tl.constexpr,
-    V: tl.constexpr,
-    BV: tl.constexpr,
-):
-    """
-    This kernel computes both cross entropy loss and the gradient of the input.
-    We only consider hard label + mean reduction for now.
-    Please refer to https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html for the math.
-    Args:
-        logits:
-            Pointer to logits tensor.
-        lse:
-            Pointer to logsumexp tensor.
-        target: Pointer to target tensor.
-        loss:
-            Pointer to tensor to store the loss.
-        V (int):
-            The number of columns in the input tensor.
-        total (int):
-            The number of non-ignored classes.
-        ignore_index (int):
-            The index to ignore in the target.
-        label_smoothing (float):
-            The amount of smoothing when computing the loss, where 0.0 means no smoothing.
-        reduction (str):
-            The string for the reduction to apply
-        BV (int):
-            The block size for vocab.
-    """
-    # https://github.com/triton-lang/triton/issues/1058
-    # If B*T*V is too large, i_n * stride will overflow out of int32, so we convert to int64
-    i_n = tl.program_id(0).to(tl.int64)
-    NV = tl.cdiv(V, BV)
-    # 1. Load target first because if the target is ignore_index, we can return right away
-    b_y = tl.load(target + i_n)
-    # 2. locate the start index
-    logits += i_n * V
-    if b_y == ignore_index:
-        # set all x as 0
-        for i in range(0, V, BV):
-            o_v = i + tl.arange(0, BV)
-            tl.store(logits + o_v, 0.0, mask=o_v < V)
-        return
-    # Online softmax: 2 loads + 1 store (compared with 3 loads + 1 store for the safe softmax)
-    # Refer to Algorithm 3 in the paper: https://arxiv.org/pdf/1805.02867
-    # 3. [Online softmax] first pass: compute logsumexp
-    # we did this in another kernel
-    b_l = tl.load(logits + b_y).to(tl.float32) * logit_scale
-    if logit_softcapping is not None:
-        b_t_y = tanh(b_l / logit_softcapping)
-        b_l = logit_softcapping * b_t_y
-        # Save the softcap derivative for the target position for use in step 6
-        b_softcap_deriv_y = 1.0 - b_t_y * b_t_y
-    b_lse = tl.load(lse + i_n)
-    # 4. Calculate the loss
-    # loss = lse - logits_l
-    b_loss = b_lse - b_l
-    # Label smoothing is a general case of normal cross entropy
-    # See the full derivation at https://github.com/linkedin/Liger-Kernel/pull/198#issue-2503665310
-    b_z = 0.0
-    eps = label_smoothing / V
-    # We need tl.debug_barrier() as mentioned in
-    # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/ops/cross_entropy.py#L34
-    tl.debug_barrier()
-    # 5. [Online Softmax] Second pass: compute gradients
-    # For 'mean' reduction, gradients are normalized by number of non-ignored elements
-    # dx_y = (softmax(x_y) - 1) / N
-    # dx_i = softmax(x_i) / N, i != y
-    # For label smoothing:
-    # dx_i = (softmax(x_y) - label_smoothing / V) / N, i != y
-    # dx_y = (softmax(x_y) - label_smoothing / V - (1 - label_smoothing)) / N
-    #      = dx_i - (1 - label_smoothing) / N
-    for iv in range(0, NV):
-        o_v = iv * BV + tl.arange(0, BV)
-        b_logits = tl.load(logits + o_v, mask=o_v < V, other=float('-inf')).to(tl.float32) * logit_scale
-        if logit_softcapping is not None:
-            b_t = tanh(b_logits / logit_softcapping)
-            b_capped = logit_softcapping * b_t
-        else:
-            b_capped = b_logits
-        if label_smoothing > 0:
-            # scale X beforehand to avoid overflow
-            b_z += tl.sum(tl.where(o_v < V, -eps * b_capped, 0.0))
-        b_p = (exp(b_capped - b_lse) - eps) * logit_scale
-        # d(softcap * tanh(x/softcap))/dx = 1 - tanh(x/softcap)^2
-        if logit_softcapping is not None:
-            b_p = b_p * (1.0 - b_t * b_t)
-        if reduction == "mean":
-            b_p = b_p / total
-        tl.store(logits + o_v, b_p, mask=o_v < V)
-        tl.debug_barrier()
-    # Original loss = H(q, p),  with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps
-    # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p)
-    #          = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i))
-    # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as:
-    #          = (1 - label_smoothing) * H(q, p) + (-sum(x_i * eps) + label_smoothing * (m + logd))
-    # Refer to H(q', p) in section 7 of the paper:
-    # https://arxiv.org/pdf/1512.00567
-    # pytorch:
-    # https://github.com/pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L516
-    # See full derivation at https://github.com/linkedin/Liger-Kernel/pull/198#issuecomment-2333753087
-    if label_smoothing > 0:
-        b_loss = b_loss * (1 - label_smoothing) + (b_z + label_smoothing * b_lse)
-    # 6. Specially handle the i==y case where `dx_y = (softmax(x_y) - (1 - label_smoothing) / N`
-    b_l = tl.load(logits + b_y)
-    # The correction term also needs the softcap chain rule factor
-    if logit_softcapping is not None:
-        b_sc_factor = b_softcap_deriv_y
-    else:
-        b_sc_factor = 1.0
-    # Normalize the loss by the number of non-ignored elements if reduction is "mean"
-    if reduction == 'mean':
-        b_loss = b_loss / total
-        b_l += (label_smoothing - 1) / total * logit_scale * b_sc_factor
-    else:
-        b_l += (label_smoothing - 1) * logit_scale * b_sc_factor
-    tl.store(loss + i_n, b_loss)
-    tl.store(logits + b_y, b_l)
-@triton.jit
-def elementwise_mul_kernel(
-    x,
-    g,
-    N: tl.constexpr,
-    B: tl.constexpr,
-):
-    """
-    This function multiplies each element of the tensor pointed by x with the value pointed by g.
-    The multiplication is performed in-place on the tensor pointed by x.
-    Parameters:
-    x:
-        Pointer to the input tensor.
-    g:
-        Pointer to the gradient output value.
-    N (int):
-        The number of columns in the input tensor.
-    B (int):
-        The block size for Triton operations.
-    """
-    # Get the program ID and convert it to int64 to avoid overflow
-    i_x = tl.program_id(0).to(tl.int64)
-    o_x = i_x * B + tl.arange(0, B)
-    # Load the gradient output value
-    b_g = tl.load(g)
-    b_x = tl.load(x + o_x, mask=o_x < N)
-    tl.store(x + o_x, b_x * b_g, mask=o_x < N)
-@dispatch('modules')
-def fused_linear_cross_entropy_forward(
     x: torch.Tensor,
     target: torch.LongTensor,
     weight: torch.Tensor,
-    bias: torch.Tensor = None,
     ignore_index: int = -100,
     label_smoothing: float = 0.0,
     logit_scale: float = 1.0,
-    logit_softcapping: float = None,
     num_chunks: int = 8,
     reduction: str = "mean",
     use_l2warp: bool = False,
     l2_penalty_factor: float = 1e-4,
     accumulate_grad_in_fp32: bool = True,
 ):
     device = x.device
-    # inputs have shape: [N, H]
-    # materialized activations will have shape: [N, V]
-    # the increase in memory = [N, V]
-    # reduction can be achieved by partitioning the number of tokens N into smaller chunks.
-    # ideally, we would like to achieve the same memory consumption as [N, H],
-    # so the expected chunk size should be:
-    # NC = ceil(V / H)
-    # C = ceil(N / NC)
-    # for ex: N = 4096*4, V = 32000, H = 4096 ==> NC = 8, C = ceil(N / NC) = 2048
     N, H, V = *x.shape, weight.shape[0]
-    BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-    # TODO: in real cases, we may need to limit the number of chunks NC to
-    # ensure the precisions of accumulated gradients
     NC = min(num_chunks, triton.cdiv(V, H))
     C = triton.next_power_of_2(triton.cdiv(N, NC))
     NC = triton.cdiv(N, C)
@@ -322,87 +146,83 @@ def fused_linear_cross_entropy_forward(
     # [N]
     loss = torch.zeros(N, device=device, dtype=torch.float)
-    total = target.ne(ignore_index).sum().item()
     for ic in range(NC):
         start, end = ic * C, min((ic + 1) * C, N)
-        # [C, N]
         c_x = x[start:end]
-        # when doing matmul, use the original precision
         # [C, V]
         c_logits = F.linear(c_x, weight, bias)
         if weight is not None and c_x.dtype != grad_dtype:
             c_x = c_x.to(dtype=grad_dtype)
         c_target = target[start:end]
-        # [C]
-        # keep lse in fp32 to maintain precision
-        c_lse = logsumexp_fwd(c_logits, scale=logit_scale, softcapping=logit_softcapping, dtype=torch.float)
-        # unreduced loss
-        c_loss = loss[start:end]
         if use_l2warp:
             c_maxx, c_ids = torch.max(c_logits, -1, keepdim=True)
-        # Here we calculate the gradient of c_logits in place so we can save memory.
-        cross_entropy_kernel[(c_logits.shape[0],)](
             logits=c_logits,
-            lse=c_lse,
             target=c_target,
-            loss=c_loss,
-            total=total,
-            ignore_index=ignore_index,
             label_smoothing=label_smoothing,
             logit_scale=logit_scale,
             logit_softcapping=logit_softcapping,
-            reduction=reduction,
-            V=V,
-            BV=BV,
-            num_warps=STATIC_WARPS,
         )
         if use_l2warp:
-            # a. Calculate the L2 gradient w.r.t logits (g_logits_l2)
             g_logits_l2 = torch.zeros_like(c_logits)
-            # Match L2Wrap: normalize by the full number of input tokens, not by non-ignored labels.
             l2_factor = l2_penalty_factor / N
             penalty_grad = c_maxx * l2_factor
             g_logits_l2.scatter_(-1, c_ids, penalty_grad)
-            # b. Backpropagate g_logits_l2 to get its effect on dx, dw, db
-            # and add it to the main gradients.
-            # Total_dx = CE_dx + L2_dx
-            # Total_dw = CE_dw + L2_dw
-            # Total_db = CE_db + L2_db
             if weight is not None:
-                torch.addmm(
-                    input=dw,
-                    mat1=g_logits_l2.t().to(dtype=grad_dtype),
-                    mat2=c_x,
-                    out=dw,
-                )
             if bias is not None:
                 torch.add(input=db, other=g_logits_l2.sum(0, dtype=bias_grad_dtype), out=db)
-            # The dx contribution must be added to the final dx calculation
             dx_l2_contribution = torch.mm(g_logits_l2, weight)
         else:
             dx_l2_contribution = 0.0
-        # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V
-        # thus dx should be of shape: C x H
         dx[start:end] = torch.mm(c_logits, weight) + dx_l2_contribution
         if weight is not None:
-            torch.addmm(
-                input=dw,
-                mat1=c_logits.t().to(dtype=grad_dtype),
-                mat2=c_x,
-                out=dw,
-            )
         if bias is not None:
             torch.add(input=db, other=c_logits.sum(0, dtype=bias_grad_dtype), out=db)
     loss = loss.sum()
     if dw is not None:
         dw = dw.to(weight)
     if db is not None:
@@ -410,51 +230,30 @@ def fused_linear_cross_entropy_forward(
     return loss, dx, dw, db
-@dispatch('modules')
-def fused_linear_cross_entropy_backward(
-    do: torch.Tensor,
-    dx: torch.Tensor,
-    dw: torch.Tensor,
-    db: torch.Tensor,
-):
-    # If cross entropy is the last layer, do is 1.0. Skip the mul to save time
-    if torch.ne(do, torch.tensor(1.0, device=do.device)):
-        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
-        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
-        N, H = dx.shape
-        B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
-        elementwise_mul_kernel[(triton.cdiv(N * H, B),)](
-            x=dx,
-            g=do,
-            N=N*H,
-            B=B,
-            num_warps=STATIC_WARPS,
-        )
-        # handle dw
-        if dw is not None:
-            V, H = dw.shape
-            elementwise_mul_kernel[(triton.cdiv(V * H, B),)](
-                x=dw,
-                g=do,
-                N=V*H,
-                B=B,
-                num_warps=STATIC_WARPS,
-            )
-        if db is not None:
-            V = db.shape[0]
-            elementwise_mul_kernel[(triton.cdiv(V, B),)](
-                x=db,
-                g=do,
-                N=V,
-                B=B,
-                num_warps=STATIC_WARPS,
-            )
     return dx, dw, db
 class FusedLinearCrossEntropyFunction(torch.autograd.Function):
     @staticmethod
@@ -464,76 +263,35 @@ class FusedLinearCrossEntropyFunction(torch.autograd.Function):
         x: torch.Tensor,
         target: torch.LongTensor,
         weight: torch.Tensor,
-        bias: torch.Tensor = None,
         ignore_index: int = -100,
         label_smoothing: float = 0.0,
         logit_scale: float = 1.0,
-        logit_softcapping: float = None,
         num_chunks: int = 8,
         reduction: str = "mean",
         use_l2warp: bool = False,
         l2_penalty_factor: float = 1e-4,
         accumulate_grad_in_fp32: bool = True,
     ):
-        """
-        Fusing the last linear layer with cross-entropy loss
-            Reference: https://github.com/mgmalek/efficient_cross_entropy
-        Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding
-        the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can
-        compute the gradient at the forward pass. By doing so, we don't have to store the x and target
-        for the backward pass.
-        x (torch.Tensor): [batch_size * seq_len, hidden_size]
-        target (torch.LongTensor): [batch_size * seq_len]
-            where each value is in [0, vocab_size).
-        weight (torch.Tensor): [vocab_size, hidden_size]
-            where `vocab_size` is the number of classes.
-        bias (Optional[torch.Tensor]): [vocab_size]
-            where `vocab_size` is the number of classes.
-        ignore_index:
-            the index to ignore in the target.
-        label_smoothing:
-            the amount of smoothing when computing the loss, where 0.0 means no smoothing.
-        logit_scale: float = 1.0,
-            A scaling factor applied to the logits. Default: 1.0
-        logit_softcapping: float = None,
-            If > 0, apply logit softcapping: logits = softcap * tanh(logits / softcap).
-            Default: 0.0
-        num_chunks: int
-            The number of chunks to split the input tensor into for processing.
-            This can help optimize memory usage and computation speed.
-            Default: 8
-        reduction:
-            Specifies the reduction to apply to the output: 'mean' | 'sum'.
-            'mean': the weighted mean of the output is taken,
-            'sum': the output will be summed.
-            Default: 'mean'.
-        use_l2warp: bool = False,
-            Whether to use L2 regularization on the logits to prevent overconfidence.
-            Default: False
-        l2_penalty_factor: float = 1e-4,
-            The L2Warp penalty factor. Default: 1e-4
-        accumulate_grad_in_fp32: bool = True,
-            Whether to accumulate weight and bias gradients in fp32 before casting them
-            back to the parameter dtype. Default: True
-        """
-        loss, dx, dw, db = fused_linear_cross_entropy_forward(
-            x,
-            target,
-            weight,
-            bias,
-            ignore_index,
-            label_smoothing,
-            logit_scale,
-            logit_softcapping,
-            num_chunks,
-            reduction,
-            use_l2warp,
-            l2_penalty_factor,
-            accumulate_grad_in_fp32,
         )
-        # downcast to dtype and store for backward
         ctx.save_for_backward(
             dx.detach(),
             dw.detach() if weight is not None else None,
@@ -545,62 +303,71 @@ class FusedLinearCrossEntropyFunction(torch.autograd.Function):
     @input_guard
     def backward(ctx, do):
         dx, dw, db = ctx.saved_tensors
-        dx, dw, db = fused_linear_cross_entropy_backward(do, dx, dw, db)
-        return dx, None, dw, db, None, None, None, None, None, None, None, None, None
 def fused_linear_cross_entropy_loss(
     x: torch.Tensor,
     target: torch.LongTensor,
     weight: torch.Tensor,
-    bias: torch.Tensor = None,
     ignore_index: int = -100,
     label_smoothing: float = 0.0,
     logit_scale: float = 1.0,
-    logit_softcapping: float = None,
     num_chunks: int = 8,
     reduction: str = "mean",
     use_l2warp: bool = False,
     l2_penalty_factor: float = 1e-4,
     accumulate_grad_in_fp32: bool = True,
-) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Args:
-        x (torch.Tensor): [batch_size * seq_len, hidden_size]
-        target (torch.LongTensor): [batch_size * seq_len]
-            where each value is in [0, vocab_size).
-        weight (torch.Tensor): [vocab_size, hidden_size]
-            where `vocab_size` is the number of classes.
-        bias (Optional[torch.Tensor]): [vocab_size]
-            where `vocab_size` is the number of classes.
-        ignore_index: int.
-            If target == ignore_index, the loss is set to 0.0.
-        label_smoothing: float
-        logit_scale: float
-            A scaling factor applied to the logits. Default: 1.0
-        logit_softcapping: float
-            If > 0, apply logit softcapping: logits = softcap * tanh(logits / softcap).
-            Default: 0.0
-        num_chunks: int
             The number of chunks to split the input tensor into for processing.
             This can help optimize memory usage and computation speed.
-            Default: 8
-        reduction:
             Specifies the reduction to apply to the output: 'mean' | 'sum'.
             'mean': the weighted mean of the output is taken,
             'sum': the output will be summed.
             Default: 'mean'.
-        use_l2warp:
-            Whether to add the L2Warp logit regularization gradient. The penalty is normalized by
-            the full number of input tokens, matching `fla.modules.l2warp.l2_warp`.
             Default: `False`.
-        l2_penalty_factor:
             The L2Warp penalty factor. Default: `1e-4`.
-        accumulate_grad_in_fp32:
-            Whether to accumulate weight and bias gradients in fp32 before casting them
-            back to the parameter dtype. Default: `True`.
     Returns:
-        losses: [batch,], float
     """
     return FusedLinearCrossEntropyFunction.apply(
         x,
@@ -616,6 +383,7 @@ def fused_linear_cross_entropy_loss(
         use_l2warp,
         l2_penalty_factor,
         accumulate_grad_in_fp32,
     )
@@ -626,41 +394,47 @@ class FusedLinearCrossEntropyLoss(nn.Module):
         ignore_index: int = -100,
         label_smoothing: float = 0.0,
         logit_scale: float = 1.0,
-        logit_softcapping: float = None,
         num_chunks: int = 8,
         reduction: str = "mean",
         use_l2warp: bool = False,
         l2_penalty_factor: float = 1e-4,
         accumulate_grad_in_fp32: bool = True,
     ):
         """
         Args:
-            ignore_index: int.
-                If target == ignore_index, the loss is set to 0.0.
-            label_smoothing: float
-            logit_scale: float
-                A scaling factor applied to the logits. Default: 1.0
-            logit_softcapping: float
-                If > 0, apply logit softcapping: logits = softcap * tanh(logits / softcap).
-                Default: 0.0
-            num_chunks: int
                 The number of chunks to split the input tensor into for processing.
                 This can help optimize memory usage and computation speed.
-                Default: 8
-            reduction:
                 Specifies the reduction to apply to the output: 'mean' | 'sum'.
                 'mean': the weighted mean of the output is taken,
                 'sum': the output will be summed.
                 Default: 'mean'.
-            use_l2warp:
-                Whether to add the L2Warp logit regularization gradient. The penalty is normalized by
-                the full number of input tokens, matching `fla.modules.l2warp.l2_warp`.
                 Default: `False`.
-            l2_penalty_factor:
                 The L2Warp penalty factor. Default: `1e-4`.
-            accumulate_grad_in_fp32:
-                Whether to accumulate weight and bias gradients in fp32 before casting them
-                back to the parameter dtype. Default: `True`.
         """
         super().__init__()
@@ -675,30 +449,28 @@ class FusedLinearCrossEntropyLoss(nn.Module):
         self.use_l2warp = use_l2warp
         self.l2_penalty_factor = l2_penalty_factor
         self.accumulate_grad_in_fp32 = accumulate_grad_in_fp32
     @torch.compiler.disable
-    def forward(
-        self,
-        x: torch.Tensor,
-        target: torch.LongTensor,
-        weight: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ):
         """
         Args:
-            x (torch.Tensor): [batch_size, seq_len, hidden_size]
-            target (torch.LongTensor): [batch_size, seq_len]
-                where each value is in [0, V).
-            weight (torch.Tensor): [vocab_size, hidden_size]
-                where `vocab_size` is the number of classes.
-            bias (Optional[torch.Tensor]): [vocab_size]
-                where `vocab_size` is the number of classes.
         Returns:
-            loss
         """
         loss = fused_linear_cross_entropy_loss(
-            x.view(-1, x.shape[-1]),
-            target.view(-1),
             weight=weight,
             bias=bias,
             ignore_index=self.ignore_index,
@@ -710,17 +482,13 @@ class FusedLinearCrossEntropyLoss(nn.Module):
             use_l2warp=self.use_l2warp,
             l2_penalty_factor=self.l2_penalty_factor,
             accumulate_grad_in_fp32=self.accumulate_grad_in_fp32,
         )
         return loss
 class LinearLossParallel(ParallelStyle):
-    def __init__(
-        self,
-        *,
-        sequence_dim: int = 1,
-        use_local_output: bool = False,
-    ):
         super().__init__()
         self.sequence_sharding = (Shard(sequence_dim),)
@@ -755,7 +523,7 @@ class LinearLossParallel(ParallelStyle):
     @staticmethod
     def _prepare_output_fn(use_local_output, mod, outputs, device_mesh):
-        return outputs.to_local() if use_local_output else outputs
     def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
         return distribute_module(

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+# adapted from https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py
+from __future__ import annotations
 from functools import partial
+from typing import TYPE_CHECKING
 import torch
 import torch.nn as nn
 from torch.distributed.tensor import Replicate, Shard, distribute_module
 from torch.distributed.tensor.parallel import ParallelStyle
+from ...backends import dispatch
+from ...modules.fused_cross_entropy.ops import cross_entropy_bwd, cross_entropy_fwd
+from ...ops.utils.op import exp, log, tanh
+from ...utils import IS_AMD, input_guard
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
 try:
     from torch.distributed.tensor import DTensor
 except (ImportError, AttributeError):
     DTensor = None
+# the hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576
 # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
+# however, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
+# the optimal maximum block size depends on your hardware, your kernel, and your dtype
 MAX_FUSED_SIZE = 65536 // 2
 STATIC_WARPS = 32 if not IS_AMD else 16
 })
 @triton.jit
 def logsumexp_fwd_kernel(
+    x,  # [N, D]
+    z,  # [N, ND]
     scale,
     softcapping,
     D: tl.constexpr,
+    BD: tl.constexpr,
     HAS_SCALE: tl.constexpr,
     HAS_SOFTCAPPING: tl.constexpr,
 ):
     i_n, i_d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    o_d = i_d * BD + tl.arange(0, BD)
     m_d = o_d < D
+    # [BD]
     b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
     if HAS_SCALE:
         b_x = b_x * scale
     if HAS_SOFTCAPPING:
         b_x = softcapping * tanh(b_x / softcapping)
+    b_x = tl.where(m_d, b_x, -float('inf'))
     b_m = tl.max(b_x, 0)
     b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m
+    tl.store(z + i_n * tl.cdiv(D, BD) + i_d, b_z)
+@triton.jit
+def elementwise_mul_kernel(
+    x,  # [N]
+    g,  # scalar
+    N: tl.constexpr,
+    BN: tl.constexpr,
 ):
+    i_x = tl.program_id(0).to(tl.int64)
+    o_x = i_x * BN + tl.arange(0, BN)
+    b_g = tl.load(g)
+    if b_g == 1.0:
+        return
+    # [BN]
+    b_x = tl.load(x + o_x, mask=o_x < N)
+    tl.store(x + o_x, b_x * b_g, mask=o_x < N)
+@dispatch
+def logsumexp_fwd(x, scale: float | None = None, softcapping: float | None = None, dtype: torch.dtype | None = None):
     shape = x.shape
     x = x.view(-1, shape[-1])
     N, D = x.shape
+    BD = min(triton.next_power_of_2(D), 64 * 1024)
+    ND = triton.cdiv(D, BD)
     z = x.new_empty(N, ND, dtype=torch.float)
+    logsumexp_fwd_kernel[(N, ND)](x=x, z=z, scale=scale, softcapping=softcapping, D=D, BD=BD)
     z = z.logsumexp(-1).view(*shape[:-1])
     if dtype is not None and dtype != torch.float:
         z = z.to(dtype)
     return z
+@dispatch
+def fused_linear_cross_entropy_fwd(
     x: torch.Tensor,
     target: torch.LongTensor,
     weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
     ignore_index: int = -100,
     label_smoothing: float = 0.0,
     logit_scale: float = 1.0,
+    logit_softcapping: float | None = None,
     num_chunks: int = 8,
     reduction: str = "mean",
     use_l2warp: bool = False,
     l2_penalty_factor: float = 1e-4,
     accumulate_grad_in_fp32: bool = True,
+    process_group: ProcessGroup | None = None,
 ):
     device = x.device
+    world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group)
+    # chunking keeps temporary logits storage close to the input tensor size.
     N, H, V = *x.shape, weight.shape[0]
     NC = min(num_chunks, triton.cdiv(V, H))
     C = triton.next_power_of_2(triton.cdiv(N, NC))
     NC = triton.cdiv(N, C)
     # [N]
     loss = torch.zeros(N, device=device, dtype=torch.float)
+    total = target.ne(ignore_index).sum().clamp_min(1)
+    dloss = total.float().reciprocal() if reduction == "mean" else x.new_ones((), dtype=torch.float)
     for ic in range(NC):
         start, end = ic * C, min((ic + 1) * C, N)
+        # [C, H]
         c_x = x[start:end]
         # [C, V]
         c_logits = F.linear(c_x, weight, bias)
         if weight is not None and c_x.dtype != grad_dtype:
             c_x = c_x.to(dtype=grad_dtype)
         c_target = target[start:end]
+        c_loss, _, c_lse, total_classes, class_start_idx = cross_entropy_fwd(
+            logits=c_logits,
+            target=c_target,
+            label_smoothing=label_smoothing,
+            logit_scale=logit_scale,
+            logit_softcapping=logit_softcapping,
+            ignore_index=ignore_index,
+            process_group=process_group,
+        )
+        loss[start:end] = c_loss
         if use_l2warp:
             c_maxx, c_ids = torch.max(c_logits, -1, keepdim=True)
+            if world_size > 1:
+                # only the shard owning the first global maximum contributes the L2 gradient
+                local_max = c_maxx.clone()
+                torch.distributed.all_reduce(c_maxx, op=torch.distributed.ReduceOp.MAX, group=process_group)
+                c_ids = torch.where(local_max == c_maxx, c_ids + class_start_idx, total_classes)
+                torch.distributed.all_reduce(c_ids, op=torch.distributed.ReduceOp.MIN, group=process_group)
+                c_ids -= class_start_idx
+                c_maxx = torch.where((c_ids >= 0) & (c_ids < V), c_maxx, 0)
+                c_ids = c_ids.clamp(0, V - 1)
+        cross_entropy_bwd(
             logits=c_logits,
             target=c_target,
+            lse=c_lse,
+            dloss=dloss.expand(end - start),
             label_smoothing=label_smoothing,
             logit_scale=logit_scale,
             logit_softcapping=logit_softcapping,
+            ignore_index=ignore_index,
+            total_classes=total_classes,
+            class_start_idx=class_start_idx,
+            inplace_backward=True,
         )
         if use_l2warp:
             g_logits_l2 = torch.zeros_like(c_logits)
+            # match L2Wrap: normalize by the full number of input tokens, not by non-ignored labels.
             l2_factor = l2_penalty_factor / N
             penalty_grad = c_maxx * l2_factor
             g_logits_l2.scatter_(-1, c_ids, penalty_grad)
             if weight is not None:
+                torch.addmm(input=dw, mat1=g_logits_l2.t().to(dtype=grad_dtype), mat2=c_x, out=dw)
             if bias is not None:
                 torch.add(input=db, other=g_logits_l2.sum(0, dtype=bias_grad_dtype), out=db)
+            # the dx contribution must be added to the final dx calculation
             dx_l2_contribution = torch.mm(g_logits_l2, weight)
         else:
             dx_l2_contribution = 0.0
         dx[start:end] = torch.mm(c_logits, weight) + dx_l2_contribution
         if weight is not None:
+            torch.addmm(input=dw, mat1=c_logits.t().to(dtype=grad_dtype), mat2=c_x, out=dw)
         if bias is not None:
             torch.add(input=db, other=c_logits.sum(0, dtype=bias_grad_dtype), out=db)
+    if world_size > 1:
+        torch.distributed.all_reduce(dx, group=process_group)
     loss = loss.sum()
+    if reduction == "mean":
+        loss = loss / total
     if dw is not None:
         dw = dw.to(weight)
     if db is not None:
     return loss, dx, dw, db
+@dispatch
+def fused_linear_cross_entropy_bwd(do: torch.Tensor, dx: torch.Tensor, dw: torch.Tensor, db: torch.Tensor | None):
+    # in-place gradient storage across repeated backward calls causes PyTorch anomalies; use Triton.
+    N, H = dx.shape
+    BN = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
+    elementwise_mul_kernel[(triton.cdiv(N * H, BN),)](x=dx, g=do, N=N*H, BN=BN, num_warps=STATIC_WARPS)
+    if dw is not None:
+        V, H = dw.shape
+        elementwise_mul_kernel[(triton.cdiv(V * H, BN),)](x=dw, g=do, N=V*H, BN=BN, num_warps=STATIC_WARPS)
+    if db is not None:
+        V = db.shape[0]
+        elementwise_mul_kernel[(triton.cdiv(V, BN),)](x=db, g=do, N=V, BN=BN, num_warps=STATIC_WARPS)
     return dx, dw, db
+fused_linear_cross_entropy_forward = fused_linear_cross_entropy_fwd
+fused_linear_cross_entropy_backward = fused_linear_cross_entropy_bwd
 class FusedLinearCrossEntropyFunction(torch.autograd.Function):
     @staticmethod
         x: torch.Tensor,
         target: torch.LongTensor,
         weight: torch.Tensor,
+        bias: torch.Tensor | None = None,
         ignore_index: int = -100,
         label_smoothing: float = 0.0,
         logit_scale: float = 1.0,
+        logit_softcapping: float | None = None,
         num_chunks: int = 8,
         reduction: str = "mean",
         use_l2warp: bool = False,
         l2_penalty_factor: float = 1e-4,
         accumulate_grad_in_fp32: bool = True,
+        process_group: ProcessGroup | None = None,
     ):
+        """Precompute gradients per token chunk so backward does not retain the logits."""
+        loss, dx, dw, db = fused_linear_cross_entropy_fwd(
+            x=x,
+            target=target,
+            weight=weight,
+            bias=bias,
+            ignore_index=ignore_index,
+            label_smoothing=label_smoothing,
+            logit_scale=logit_scale,
+            logit_softcapping=logit_softcapping,
+            num_chunks=num_chunks,
+            reduction=reduction,
+            use_l2warp=use_l2warp,
+            l2_penalty_factor=l2_penalty_factor,
+            accumulate_grad_in_fp32=accumulate_grad_in_fp32,
+            process_group=process_group,
         )
         ctx.save_for_backward(
             dx.detach(),
             dw.detach() if weight is not None else None,
     @input_guard
     def backward(ctx, do):
         dx, dw, db = ctx.saved_tensors
+        dx, dw, db = fused_linear_cross_entropy_bwd(do=do, dx=dx, dw=dw, db=db)
+        return dx, None, dw, db, None, None, None, None, None, None, None, None, None, None
 def fused_linear_cross_entropy_loss(
     x: torch.Tensor,
     target: torch.LongTensor,
     weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
     ignore_index: int = -100,
     label_smoothing: float = 0.0,
     logit_scale: float = 1.0,
+    logit_softcapping: float | None = None,
     num_chunks: int = 8,
     reduction: str = "mean",
     use_l2warp: bool = False,
     l2_penalty_factor: float = 1e-4,
     accumulate_grad_in_fp32: bool = True,
+    process_group: ProcessGroup | None = None,
+) -> torch.Tensor:
     """
     Args:
+        x (torch.Tensor):
+            Input of shape `[batch_size * seq_len, hidden_size]`.
+        target (torch.LongTensor):
+            Target indices of shape `[batch_size * seq_len]` with values in `[0, vocab_size)`.
+        weight (torch.Tensor):
+            Projection weight of shape `[vocab_size, hidden_size]` where `vocab_size` is the number of classes.
+        bias (torch.Tensor, Optional):
+            Projection bias of shape `[vocab_size]`. Default: `None`.
+        ignore_index (int, Optional):
+            Target index excluded from the loss. Default: -100.
+        label_smoothing (float, Optional):
+            Uniform label smoothing coefficient. Default: 0.0.
+        logit_scale (float, Optional):
+            A scaling factor applied to the logits. Default: 1.0.
+        logit_softcapping (float, Optional):
+            If not `None`, apply logit softcapping: logits = softcap * tanh(logits / softcap).
+            Default: `None`.
+        num_chunks (int, Optional):
             The number of chunks to split the input tensor into for processing.
             This can help optimize memory usage and computation speed.
+            Default: 8.
+        reduction (str, Optional):
             Specifies the reduction to apply to the output: 'mean' | 'sum'.
             'mean': the weighted mean of the output is taken,
             'sum': the output will be summed.
             Default: 'mean'.
+        use_l2warp (bool, Optional):
+            Whether to add the L2Warp logit regularization gradient.
+            The penalty is normalized by the full number of input tokens, matching `fla.modules.l2warp.l2_warp`.
             Default: `False`.
+        l2_penalty_factor (float, Optional):
             The L2Warp penalty factor. Default: `1e-4`.
+        accumulate_grad_in_fp32 (bool, Optional):
+            Whether to accumulate weight and bias gradients in fp32 before casting them back to the parameter dtype.
+            Default: `True`.
+        process_group (ProcessGroup, Optional):
+            Group with equal contiguous vocabulary shards of weight and bias, and replicated inputs and targets.
+            Loss and input gradients are replicated; weight and bias gradients remain sharded. Default: `None`.
+            Multi-rank groups are not supported by the Ascend backend.
     Returns:
+        torch.Tensor:
+            Scalar loss in fp32.
     """
     return FusedLinearCrossEntropyFunction.apply(
         x,
         use_l2warp,
         l2_penalty_factor,
         accumulate_grad_in_fp32,
+        process_group,
     )
         ignore_index: int = -100,
         label_smoothing: float = 0.0,
         logit_scale: float = 1.0,
+        logit_softcapping: float | None = None,
         num_chunks: int = 8,
         reduction: str = "mean",
         use_l2warp: bool = False,
         l2_penalty_factor: float = 1e-4,
         accumulate_grad_in_fp32: bool = True,
+        process_group: ProcessGroup | None = None,
     ):
         """
         Args:
+            ignore_index (int, Optional):
+                Target index excluded from the loss. Default: -100.
+            label_smoothing (float, Optional):
+                Uniform label smoothing coefficient. Default: 0.0.
+            logit_scale (float, Optional):
+                A scaling factor applied to the logits. Default: 1.0.
+            logit_softcapping (float, Optional):
+                If not `None`, apply logit softcapping: logits = softcap * tanh(logits / softcap).
+                Default: `None`.
+            num_chunks (int, Optional):
                 The number of chunks to split the input tensor into for processing.
                 This can help optimize memory usage and computation speed.
+                Default: 8.
+            reduction (str, Optional):
                 Specifies the reduction to apply to the output: 'mean' | 'sum'.
                 'mean': the weighted mean of the output is taken,
                 'sum': the output will be summed.
                 Default: 'mean'.
+            use_l2warp (bool, Optional):
+                Whether to add the L2Warp logit regularization gradient.
+                The penalty is normalized by the full number of input tokens, matching `fla.modules.l2warp.l2_warp`.
                 Default: `False`.
+            l2_penalty_factor (float, Optional):
                 The L2Warp penalty factor. Default: `1e-4`.
+            accumulate_grad_in_fp32 (bool, Optional):
+                Whether to accumulate weight and bias gradients in fp32 before casting them back to the parameter dtype.
+                Default: `True`.
+            process_group (ProcessGroup, Optional):
+                Group with equal contiguous vocabulary shards of weight and bias, and replicated inputs and targets.
+                Loss and input gradients are replicated; weight and bias gradients remain sharded. Default: `None`.
+                Multi-rank groups are not supported by the Ascend backend.
         """
         super().__init__()
         self.use_l2warp = use_l2warp
         self.l2_penalty_factor = l2_penalty_factor
         self.accumulate_grad_in_fp32 = accumulate_grad_in_fp32
+        self.process_group = process_group
     @torch.compiler.disable
+    def forward(self, x: torch.Tensor, target: torch.LongTensor, weight: torch.Tensor, bias: torch.Tensor | None = None):
         """
         Args:
+            x (torch.Tensor):
+                Input of shape `[batch_size, seq_len, hidden_size]`.
+            target (torch.LongTensor):
+                Target indices of shape `[batch_size, seq_len]` with values in `[0, vocab_size)`.
+            weight (torch.Tensor):
+                Projection weight of shape `[vocab_size, hidden_size]` where `vocab_size` is the number of classes.
+            bias (torch.Tensor, Optional):
+                Projection bias of shape `[vocab_size]`. Default: `None`.
         Returns:
+            torch.Tensor:
+                Scalar loss in fp32.
         """
         loss = fused_linear_cross_entropy_loss(
+            x=x.reshape(-1, x.shape[-1]),
+            target=target.reshape(-1),
             weight=weight,
             bias=bias,
             ignore_index=self.ignore_index,
             use_l2warp=self.use_l2warp,
             l2_penalty_factor=self.l2_penalty_factor,
             accumulate_grad_in_fp32=self.accumulate_grad_in_fp32,
+            process_group=self.process_group,
         )
         return loss
 class LinearLossParallel(ParallelStyle):
+    def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False):
         super().__init__()
         self.sequence_sharding = (Shard(sequence_dim),)
     @staticmethod
     def _prepare_output_fn(use_local_output, mod, outputs, device_mesh):
+        return outputs.to_local() if use_local_output and isinstance(outputs, DTensor) else outputs
     def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
         return distribute_module(

build/torch-cuda/modules/fused_norm_gate.py DELETED Viewed

@@ -1,1245 +0,0 @@
-# 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 __future__ import annotations
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import triton
-import triton.language as tl
-from ..utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard
-@triton.heuristics(
-    {
-        "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
-        "HAS_RESIDUAL": lambda args: args["residual"] is not None,
-        "HAS_WEIGHT": lambda args: args["w"] is not None,
-        "HAS_BIAS": lambda args: args["b"] is not None,
-    }
-)
-@triton.autotune(
-    configs=[triton.Config({"BT": BT}, num_warps=num_warps) for BT in [16, 32, 64] for num_warps in [4, 8, 16]],
-    key=["D", "NB", "IS_RMS_NORM", "STORE_RESIDUAL_OUT", "HAS_RESIDUAL", "HAS_WEIGHT"],
-    **autotune_cache_kwargs,
-)
-@triton.jit
-def layer_norm_gated_fwd_kernel(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    y,  # pointer to the output
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    residual,  # pointer to the residual
-    residual_out,  # pointer to the residual
-    mean,  # pointer to the mean
-    rstd,  # pointer to the 1/std
-    eps,  # epsilon to avoid division by zero
-    T,  # number of rows in x
-    D: tl.constexpr,  # number of columns in x
-    BT: tl.constexpr,
-    BD: tl.constexpr,
-    NB: tl.constexpr,
-    ACTIVATION: tl.constexpr,
-    IS_RMS_NORM: tl.constexpr,
-    STORE_RESIDUAL_OUT: tl.constexpr,
-    HAS_RESIDUAL: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-):
-    i_t = tl.program_id(0)
-    o_d = tl.arange(0, BD)
-    m_d = o_d < D
-    p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-    b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
-    if HAS_RESIDUAL:
-        p_res = tl.make_block_ptr(residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-        b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32)
-    if STORE_RESIDUAL_OUT:
-        p_res_out = tl.make_block_ptr(residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-        tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1))
-    if not IS_RMS_NORM:
-        b_mean = tl.sum(b_x, axis=1) / D
-        p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,))
-        tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,))
-        b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)
-        b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
-    else:
-        b_xbar = tl.where(m_d[None, :], b_x, 0.0)
-        b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
-    b_rstd = 1 / tl.sqrt(b_var + eps)
-    p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))
-    tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))
-    if HAS_WEIGHT:
-        b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
-    if HAS_BIAS:
-        b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
-    b_x_hat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None]
-    b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat
-    if HAS_BIAS:
-        b_y = b_y + b_b[None, :]
-    # swish/sigmoid output gate
-    p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-    b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
-    if ACTIVATION == "swish" or ACTIVATION == "silu":
-        b_y = b_y * b_g * tl.sigmoid(b_g)
-    elif ACTIVATION == "sigmoid":
-        b_y = b_y * tl.sigmoid(b_g)
-    # Write output
-    p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-    tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
-@triton.heuristics(
-    {
-        "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
-        "HAS_RESIDUAL": lambda args: args["residual"] is not None,
-        "HAS_WEIGHT": lambda args: args["w"] is not None,
-        "HAS_BIAS": lambda args: args["b"] is not None,
-    }
-)
-@triton.autotune(
-    configs=[triton.Config({}, num_warps=num_warps) for num_warps in [2, 4, 8, 16]],
-    key=["D", "IS_RMS_NORM", "STORE_RESIDUAL_OUT", "HAS_RESIDUAL", "HAS_WEIGHT"],
-    **autotune_cache_kwargs,
-)
-@triton.jit
-def layer_norm_gated_fwd_kernel1(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    y,  # pointer to the output
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    residual,  # pointer to the residual
-    residual_out,  # pointer to the residual
-    mean,  # pointer to the mean
-    rstd,  # pointer to the 1/std
-    eps,  # epsilon to avoid division by zero
-    D: tl.constexpr,  # number of columns in x
-    BD: tl.constexpr,
-    ACTIVATION: tl.constexpr,
-    IS_RMS_NORM: tl.constexpr,
-    STORE_RESIDUAL_OUT: tl.constexpr,
-    HAS_RESIDUAL: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-):
-    i_t = tl.program_id(0)
-    x += i_t * D
-    y += i_t * D
-    g += i_t * D
-    if HAS_RESIDUAL:
-        residual += i_t * D
-    if STORE_RESIDUAL_OUT:
-        residual_out += i_t * D
-    o_d = tl.arange(0, BD)
-    m_d = o_d < D
-    b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32)
-    if HAS_RESIDUAL:
-        b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32)
-    if STORE_RESIDUAL_OUT:
-        tl.store(residual_out + o_d, b_x, mask=m_d)
-    if not IS_RMS_NORM:
-        b_mean = tl.sum(b_x, axis=0) / D
-        tl.store(mean + i_t, b_mean)
-        b_xbar = tl.where(m_d, b_x - b_mean, 0.0)
-        b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
-    else:
-        b_xbar = tl.where(m_d, b_x, 0.0)
-        b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
-    b_rstd = 1 / tl.sqrt(b_var + eps)
-    tl.store(rstd + i_t, b_rstd)
-    if HAS_WEIGHT:
-        b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
-    if HAS_BIAS:
-        b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
-    b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
-    b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat
-    if HAS_BIAS:
-        b_y = b_y + b_b
-    # swish/sigmoid output gate
-    b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32)
-    if ACTIVATION == "swish" or ACTIVATION == "silu":
-        b_y = b_y * b_g * tl.sigmoid(b_g)
-    elif ACTIVATION == "sigmoid":
-        b_y = b_y * tl.sigmoid(b_g)
-    # Write output
-    tl.store(y + o_d, b_y, mask=m_d)
-@triton.heuristics(
-    {
-        "HAS_DRESIDUAL": lambda args: args["dresidual"] is not None,
-        "HAS_WEIGHT": lambda args: args["w"] is not None,
-        "HAS_BIAS": lambda args: args["b"] is not None,
-        "RECOMPUTE_OUTPUT": lambda args: args["y"] is not None,
-    }
-)
-@triton.autotune(
-    configs=[triton.Config({"BT": BT}, num_warps=num_warps) for BT in [16, 32, 64] for num_warps in [4, 8, 16]],
-    key=["D", "NB", "IS_RMS_NORM", "HAS_DRESIDUAL", "HAS_WEIGHT"],
-    **autotune_cache_kwargs,
-)
-@triton.jit
-def layer_norm_gated_bwd_kernel(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    y,  # pointer to the output to be recomputed
-    dy,  # pointer to the output gradient
-    dx,  # pointer to the input gradient
-    dg,  # pointer to the gate gradient
-    dw,  # pointer to the partial sum of weights gradient
-    db,  # pointer to the partial sum of biases gradient
-    dresidual,
-    dresidual_in,
-    mean,
-    rstd,
-    T,
-    BS,
-    D: tl.constexpr,
-    BT: tl.constexpr,
-    BD: tl.constexpr,
-    NB: tl.constexpr,
-    ACTIVATION: tl.constexpr,
-    IS_RMS_NORM: tl.constexpr,
-    STORE_DRESIDUAL: tl.constexpr,
-    HAS_DRESIDUAL: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    RECOMPUTE_OUTPUT: tl.constexpr,
-):
-    i_s = tl.program_id(0)
-    o_d = tl.arange(0, BD)
-    m_d = o_d < D
-    if HAS_WEIGHT:
-        b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
-        b_dw = tl.zeros((BT, BD), dtype=tl.float32)
-    if HAS_BIAS:
-        b_b = tl.load(b + o_d, mask=m_d, other=0.0).to(tl.float32)
-        b_db = tl.zeros((BT, BD), dtype=tl.float32)
-    # the caller guarantees NS = min(SM, T), so every program has at least one token.
-    # the last program's range may slightly exceed T (since BS = ceil(T/NS));
-    # make_block_ptr uses the true tensor shape (T, D), so boundary_check
-    # handles the partial tail tile by zero-padding loads and skipping stores.
-    # the m_t mask below further ensures dw/db only accumulate valid rows (< T).
-    for i_t in range(i_s * BS, i_s * BS + BS, BT):
-        p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_dg = tl.make_block_ptr(dg, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        # [BT, BD]
-        b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
-        b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
-        b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32)
-        if not IS_RMS_NORM:
-            p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t,), (BT,), (0,))
-            b_mean = tl.load(p_mean, boundary_check=(0,))
-        p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t,), (BT,), (0,))
-        b_rstd = tl.load(p_rstd, boundary_check=(0,))
-        # Compute dx
-        b_xhat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None]
-        b_xhat = tl.where(m_d[None, :], b_xhat, 0.0)
-        b_y = b_xhat * b_w[None, :] if HAS_WEIGHT else b_xhat
-        if HAS_BIAS:
-            b_y = b_y + b_b[None, :]
-        if RECOMPUTE_OUTPUT:
-            p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-            tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
-        b_sigmoid_g = tl.sigmoid(b_g)
-        if ACTIVATION == "swish" or ACTIVATION == "silu":
-            b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
-            b_dy = b_dy * b_g * b_sigmoid_g
-        elif ACTIVATION == "sigmoid":
-            b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
-            b_dy = b_dy * b_sigmoid_g
-        b_wdy = b_dy
-        if HAS_WEIGHT or HAS_BIAS:
-            # when BT > BS, a tile may span into the next program's range;
-            # mask to this program's upper bound to avoid double-counting dw/db.
-            m_t = (i_t + tl.arange(0, BT)) < min(i_s * BS + BS, T)
-        if HAS_WEIGHT:
-            b_wdy = b_dy * b_w
-            b_dw += tl.where(m_t[:, None], b_dy * b_xhat, 0.0)
-        if HAS_BIAS:
-            b_db += tl.where(m_t[:, None], b_dy, 0.0)
-        if not IS_RMS_NORM:
-            b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D
-            b_c2 = tl.sum(b_wdy, axis=1) / D
-            b_dx = (b_wdy - (b_xhat * b_c1[:, None] + b_c2[:, None])) * b_rstd[:, None]
-        else:
-            b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D
-            b_dx = (b_wdy - b_xhat * b_c1[:, None]) * b_rstd[:, None]
-        if HAS_DRESIDUAL:
-            p_dres = tl.make_block_ptr(dresidual, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-            b_dres = tl.load(p_dres, boundary_check=(0, 1)).to(tl.float32)
-            b_dx += b_dres
-        # Write dx
-        if STORE_DRESIDUAL:
-            p_dres_in = tl.make_block_ptr(dresidual_in, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-            tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1))
-    if HAS_WEIGHT:
-        tl.store(dw + i_s * D + o_d, tl.sum(b_dw, axis=0), mask=m_d)
-    if HAS_BIAS:
-        tl.store(db + i_s * D + o_d, tl.sum(b_db, axis=0), mask=m_d)
-@triton.heuristics(
-    {
-        "HAS_DRESIDUAL": lambda args: args["dresidual"] is not None,
-        "HAS_WEIGHT": lambda args: args["w"] is not None,
-        "HAS_BIAS": lambda args: args["b"] is not None,
-        "RECOMPUTE_OUTPUT": lambda args: args["y"] is not None,
-    }
-)
-@triton.autotune(
-    configs=[triton.Config({}, num_warps=num_warps) for num_warps in [2, 4, 8, 16]],
-    key=["D", "IS_RMS_NORM", "STORE_DRESIDUAL", "HAS_DRESIDUAL", "HAS_WEIGHT"],
-    **autotune_cache_kwargs,
-)
-@triton.jit
-def layer_norm_gated_bwd_kernel1(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    y,  # pointer to the output to be recomputed
-    dy,  # pointer to the output gradient
-    dx,  # pointer to the input gradient
-    dg,  # pointer to the gate gradient
-    dw,  # pointer to the partial sum of weights gradient
-    db,  # pointer to the partial sum of biases gradient
-    dresidual,
-    dresidual_in,
-    mean,
-    rstd,
-    T,
-    BS,
-    D: tl.constexpr,
-    BD: tl.constexpr,
-    ACTIVATION: tl.constexpr,
-    IS_RMS_NORM: tl.constexpr,
-    STORE_DRESIDUAL: tl.constexpr,
-    HAS_DRESIDUAL: tl.constexpr,
-    HAS_WEIGHT: tl.constexpr,
-    HAS_BIAS: tl.constexpr,
-    RECOMPUTE_OUTPUT: tl.constexpr,
-):
-    i_s = tl.program_id(0)
-    o_d = tl.arange(0, BD)
-    mask = o_d < D
-    x += i_s * BS * D
-    g += i_s * BS * D
-    if HAS_DRESIDUAL:
-        dresidual += i_s * BS * D
-    if STORE_DRESIDUAL:
-        dresidual_in += i_s * BS * D
-    dy += i_s * BS * D
-    dx += i_s * BS * D
-    dg += i_s * BS * D
-    if RECOMPUTE_OUTPUT:
-        y += i_s * BS * D
-    if HAS_WEIGHT:
-        b_w = tl.load(w + o_d, mask=mask).to(tl.float32)
-        b_dw = tl.zeros((BD,), dtype=tl.float32)
-    if HAS_BIAS:
-        b_b = tl.load(b + o_d, mask=mask, other=0.0).to(tl.float32)
-        b_db = tl.zeros((BD,), dtype=tl.float32)
-    for i_t in range(i_s * BS, min(i_s * BS + BS, T)):
-        # Load data to SRAM
-        b_x = tl.load(x + o_d, mask=mask, other=0).to(tl.float32)
-        b_g = tl.load(g + o_d, mask=mask, other=0).to(tl.float32)
-        b_dy = tl.load(dy + o_d, mask=mask, other=0).to(tl.float32)
-        if not IS_RMS_NORM:
-            b_mean = tl.load(mean + i_t)
-        b_rstd = tl.load(rstd + i_t)
-        # Compute dx
-        b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
-        b_xhat = tl.where(mask, b_xhat, 0.0)
-        b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat
-        if HAS_BIAS:
-            b_y = b_y + b_b
-        if RECOMPUTE_OUTPUT:
-            tl.store(y + o_d, b_y, mask=mask)
-        b_sigmoid_g = tl.sigmoid(b_g)
-        if ACTIVATION == "swish" or ACTIVATION == "silu":
-            b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
-            b_dy = b_dy * b_g * b_sigmoid_g
-        elif ACTIVATION == "sigmoid":
-            b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
-            b_dy = b_dy * b_sigmoid_g
-        b_wdy = b_dy
-        if HAS_WEIGHT:
-            b_wdy = b_dy * b_w
-            b_dw += b_dy * b_xhat
-        if HAS_BIAS:
-            b_db += b_dy
-        if not IS_RMS_NORM:
-            b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D
-            b_c2 = tl.sum(b_wdy, axis=0) / D
-            b_dx = (b_wdy - (b_xhat * b_c1 + b_c2)) * b_rstd
-        else:
-            b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D
-            b_dx = (b_wdy - b_xhat * b_c1) * b_rstd
-        if HAS_DRESIDUAL:
-            b_dres = tl.load(dresidual + o_d, mask=mask, other=0).to(tl.float32)
-            b_dx += b_dres
-        # Write dx
-        if STORE_DRESIDUAL:
-            tl.store(dresidual_in + o_d, b_dx, mask=mask)
-        tl.store(dx + o_d, b_dx, mask=mask)
-        tl.store(dg + o_d, b_dg, mask=mask)
-        x += D
-        g += D
-        if HAS_DRESIDUAL:
-            dresidual += D
-        if STORE_DRESIDUAL:
-            dresidual_in += D
-        if RECOMPUTE_OUTPUT:
-            y += D
-        dy += D
-        dx += D
-        dg += D
-    if HAS_WEIGHT:
-        tl.store(dw + i_s * D + o_d, b_dw, mask=mask)
-    if HAS_BIAS:
-        tl.store(db + i_s * D + o_d, b_db, mask=mask)
-def layer_norm_gated_fwd(
-    x: torch.Tensor,
-    g: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    activation: str = "swish",
-    eps: float = 1e-5,
-    residual: torch.Tensor = None,
-    out_dtype: torch.dtype = None,
-    residual_dtype: torch.dtype = None,
-    is_rms_norm: bool = False,
-):
-    if residual is not None:
-        residual_dtype = residual.dtype
-    T, D = x.shape
-    if residual is not None:
-        assert residual.shape == (T, D)
-    if weight is not None:
-        assert weight.shape == (D,)
-    if bias is not None:
-        assert bias.shape == (D,)
-    # allocate output
-    y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
-    if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype):
-        residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
-    else:
-        residual_out = None
-    mean = torch.empty((T,), dtype=torch.float, device=x.device) if not is_rms_norm else None
-    rstd = torch.empty((T,), dtype=torch.float, device=x.device)
-    # Less than 64KB per feature: enqueue fused kernel
-    MAX_FUSED_SIZE = 65536 // x.element_size()
-    BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
-    if D > BD:
-        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-    # heuristics for number of warps
-    if D <= 512:
-        # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range
-        # of T before recompiling the kernel.
-        # NB = triton.cdiv(T, 2048)
-        NB = triton.cdiv(T, 2048 * 32)
-        def grid(meta):
-            return (triton.cdiv(T, meta["BT"]),)
-        layer_norm_gated_fwd_kernel[grid](
-            x=x,
-            g=g,
-            y=y,
-            w=weight,
-            b=bias,
-            residual=residual,
-            residual_out=residual_out,
-            mean=mean,
-            rstd=rstd,
-            eps=eps,
-            T=T,
-            D=D,
-            BD=BD,
-            NB=NB,
-            ACTIVATION=activation,
-            IS_RMS_NORM=is_rms_norm,
-        )
-    else:
-        layer_norm_gated_fwd_kernel1[(T,)](
-            x=x,
-            g=g,
-            y=y,
-            w=weight,
-            b=bias,
-            residual=residual,
-            residual_out=residual_out,
-            mean=mean,
-            rstd=rstd,
-            eps=eps,
-            D=D,
-            BD=BD,
-            ACTIVATION=activation,
-            IS_RMS_NORM=is_rms_norm,
-        )
-    # residual_out is None if residual is None and residual_dtype == input_dtype
-    return y, mean, rstd, residual_out if residual_out is not None else x
-def layer_norm_gated_bwd(
-    dy: torch.Tensor,
-    x: torch.Tensor,
-    g: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    activation: str = "swish",
-    eps: float = 1e-5,
-    mean: torch.Tensor = None,
-    rstd: torch.Tensor = None,
-    dresidual: torch.Tensor = None,
-    has_residual: bool = False,
-    is_rms_norm: bool = False,
-    x_dtype: torch.dtype = None,
-    recompute_output: bool = False,
-):
-    T, D = x.shape
-    assert dy.shape == (T, D)
-    if dresidual is not None:
-        assert dresidual.shape == (T, D)
-    if weight is not None:
-        assert weight.shape == (D,)
-    if bias is not None:
-        assert bias.shape == (D,)
-    # allocate output
-    dx = torch.empty_like(x) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
-    dg = torch.empty_like(g) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
-    dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None
-    y = torch.empty(T, D, dtype=dy.dtype, device=dy.device) if recompute_output else None
-    # Less than 64KB per feature: enqueue fused kernel
-    MAX_FUSED_SIZE = 65536 // x.element_size()
-    BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
-    if D > BD:
-        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-    # cap program count to T so no program is completely idle.
-    # without this, high-SM GPUs (e.g. B200, 160 SMs) with small T would
-    # launch idle programs whose make_block_ptr offsets exceed the tensor shape.
-    NS = min(get_multiprocessor_count(x.device.index), T)
-    BS = math.ceil(T / NS)
-    dw = torch.empty((NS, D), dtype=torch.float, device=weight.device) if weight is not None else None
-    db = torch.empty((NS, D), dtype=torch.float, device=bias.device) if bias is not None else None
-    grid = (NS,)
-    if D <= 512:
-        # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range
-        # of T before recompiling the kernel.
-        # NB = triton.cdiv(T, 2048)
-        NB = triton.cdiv(T, 2048 * 32)
-        layer_norm_gated_bwd_kernel[grid](
-            x=x,
-            g=g,
-            w=weight,
-            b=bias,
-            y=y,
-            dy=dy,
-            dx=dx,
-            dg=dg,
-            dw=dw,
-            db=db,
-            dresidual=dresidual,
-            dresidual_in=dresidual_in,
-            mean=mean,
-            rstd=rstd,
-            T=T,
-            D=D,
-            BS=BS,
-            BD=BD,
-            NB=NB,
-            ACTIVATION=activation,
-            IS_RMS_NORM=is_rms_norm,
-            STORE_DRESIDUAL=dresidual_in is not None,
-        )
-    else:
-        layer_norm_gated_bwd_kernel1[grid](
-            x=x,
-            g=g,
-            w=weight,
-            b=bias,
-            y=y,
-            dy=dy,
-            dx=dx,
-            dg=dg,
-            dw=dw,
-            db=db,
-            dresidual=dresidual,
-            dresidual_in=dresidual_in,
-            mean=mean,
-            rstd=rstd,
-            T=T,
-            D=D,
-            BS=BS,
-            BD=BD,
-            ACTIVATION=activation,
-            IS_RMS_NORM=is_rms_norm,
-            STORE_DRESIDUAL=dresidual_in is not None,
-        )
-    dw = dw.sum(0).to(weight.dtype) if weight is not None else None
-    db = db.sum(0).to(bias.dtype) if bias is not None else None
-    # Don't need to compute dresidual_in separately in this case
-    if has_residual and dx.dtype == x.dtype:
-        dresidual_in = dx
-    return (dx, dg, dw, db, dresidual_in) if not recompute_output else (dx, dg, dw, db, dresidual_in, y)
-class LayerNormGatedFunction(torch.autograd.Function):
-    @staticmethod
-    @input_guard
-    def forward(
-        ctx,
-        x: torch.Tensor,
-        g: torch.Tensor,
-        weight: torch.Tensor,
-        bias: torch.Tensor,
-        activation: str,
-        residual: torch.Tensor | None = None,
-        eps: float = 1e-6,
-        prenorm: bool = False,
-        residual_in_fp32: bool = False,
-        is_rms_norm: bool = False,
-    ):
-        x_shape_og = x.shape
-        g_shape_og = g.shape
-        # reshape input data into 2D tensor
-        x = x.reshape(-1, x.shape[-1])
-        g = g.reshape(-1, g.shape[-1])
-        if residual is not None:
-            assert residual.shape == x_shape_og
-            residual = residual.reshape(-1, residual.shape[-1])
-        residual_dtype = residual.dtype if residual is not None else (torch.float if residual_in_fp32 else None)
-        y, mean, rstd, residual_out = layer_norm_gated_fwd(
-            x=x,
-            g=g,
-            weight=weight,
-            bias=bias,
-            activation=activation,
-            eps=eps,
-            residual=residual,
-            residual_dtype=residual_dtype,
-            is_rms_norm=is_rms_norm,
-        )
-        ctx.save_for_backward(residual_out, g, weight, bias, mean, rstd)
-        ctx.x_shape_og = x_shape_og
-        ctx.g_shape_og = g_shape_og
-        ctx.activation = activation
-        ctx.eps = eps
-        ctx.is_rms_norm = is_rms_norm
-        ctx.has_residual = residual is not None
-        ctx.prenorm = prenorm
-        ctx.x_dtype = x.dtype
-        y = y.reshape(x_shape_og)
-        return y if not prenorm else (y, residual_out.reshape(x_shape_og))
-    @staticmethod
-    @input_guard
-    def backward(ctx, dy, *args):
-        x, g, weight, bias, mean, rstd = ctx.saved_tensors
-        dy = dy.reshape(-1, dy.shape[-1])
-        assert dy.shape == x.shape
-        if ctx.prenorm:
-            dresidual = args[0]
-            dresidual = dresidual.reshape(-1, dresidual.shape[-1])
-            assert dresidual.shape == x.shape
-        else:
-            dresidual = None
-        dx, dg, dw, db, dres_in = layer_norm_gated_bwd(
-            dy=dy,
-            x=x,
-            g=g,
-            weight=weight,
-            bias=bias,
-            activation=ctx.activation,
-            eps=ctx.eps,
-            mean=mean,
-            rstd=rstd,
-            dresidual=dresidual,
-            has_residual=ctx.has_residual,
-            is_rms_norm=ctx.is_rms_norm,
-            x_dtype=ctx.x_dtype,
-        )
-        return (
-            dx.reshape(ctx.x_shape_og),
-            dg.reshape(ctx.g_shape_og),
-            dw,
-            db,
-            None,
-            dres_in.reshape(ctx.x_shape_og) if ctx.has_residual else None,
-            None,
-            None,
-            None,
-            None,
-        )
-class LayerNormGatedLinearFunction(torch.autograd.Function):
-    @staticmethod
-    @input_guard
-    def forward(
-        ctx,
-        x: torch.Tensor,
-        g: torch.Tensor,
-        norm_weight: torch.Tensor,
-        norm_bias: torch.Tensor,
-        linear_weight: torch.Tensor,
-        linear_bias: torch.Tensor,
-        residual: torch.Tensor | None = None,
-        eps: float = 1e-6,
-        prenorm: bool = False,
-        residual_in_fp32: bool = False,
-        is_rms_norm: bool = False,
-    ):
-        x_shape_og = x.shape
-        g_shape_og = g.shape
-        # reshape input data into 2D tensor
-        x = x.reshape(-1, x.shape[-1])
-        g = g.reshape(-1, g.shape[-1])
-        if residual is not None:
-            assert residual.shape == x_shape_og
-            residual = residual.reshape(-1, residual.shape[-1])
-        residual_dtype = residual.dtype if residual is not None else (torch.float if residual_in_fp32 else None)
-        y, mean, rstd, residual_out = layer_norm_gated_fwd(
-            x=x,
-            g=g,
-            weight=norm_weight,
-            bias=norm_bias,
-            eps=eps,
-            residual=residual,
-            residual_dtype=residual_dtype,
-            is_rms_norm=is_rms_norm,
-        )
-        y = y.reshape(x_shape_og)
-        dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype
-        linear_weight = linear_weight.to(dtype)
-        linear_bias = linear_bias.to(dtype) if linear_bias is not None else None
-        out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias)
-        # We don't store y, will be recomputed in the backward pass to save memory
-        ctx.save_for_backward(residual_out, g, norm_weight, norm_bias, linear_weight, mean, rstd)
-        ctx.x_shape_og = x_shape_og
-        ctx.g_shape_og = g_shape_og
-        ctx.eps = eps
-        ctx.is_rms_norm = is_rms_norm
-        ctx.has_residual = residual is not None
-        ctx.prenorm = prenorm
-        ctx.x_dtype = x.dtype
-        ctx.linear_bias_is_none = linear_bias is None
-        return out if not prenorm else (out, residual_out.reshape(x_shape_og))
-    @staticmethod
-    @input_guard
-    def backward(ctx, dout, *args):
-        x, g, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors
-        dout = dout.reshape(-1, dout.shape[-1])
-        dy = F.linear(dout, linear_weight.t())
-        dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0)
-        assert dy.shape == x.shape
-        if ctx.prenorm:
-            dresidual = args[0]
-            dresidual = dresidual.reshape(-1, dresidual.shape[-1])
-            assert dresidual.shape == x.shape
-        else:
-            dresidual = None
-        dx, dg, dnorm_weight, dnorm_bias, dres_in, y = layer_norm_gated_bwd(
-            dy=dy,
-            x=x,
-            g=g,
-            weight=norm_weight,
-            bias=norm_bias,
-            eps=ctx.eps,
-            mean=mean,
-            rstd=rstd,
-            dresidual=dresidual,
-            has_residual=ctx.has_residual,
-            is_rms_norm=ctx.is_rms_norm,
-            x_dtype=ctx.x_dtype,
-            recompute_output=True,
-        )
-        dlinear_weight = torch.einsum("bo,bi->oi", dout, y)
-        return (
-            dx.reshape(ctx.x_shape_og),
-            dg.reshape(ctx.g_shape_og),
-            dnorm_weight,
-            dnorm_bias,
-            dlinear_weight,
-            dlinear_bias,
-            dres_in.reshape(ctx.x_shape_og) if ctx.has_residual else None,
-            None,
-            None,
-            None,
-            None,
-        )
-def layer_norm_gated(
-    x: torch.Tensor,
-    g: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    activation: str = "swish",
-    residual: torch.Tensor | None = None,
-    prenorm: bool = False,
-    residual_in_fp32: bool = False,
-    eps: float = 1e-6,
-):
-    return LayerNormGatedFunction.apply(
-        x,
-        g,
-        weight,
-        bias,
-        activation,
-        residual,
-        eps,
-        prenorm,
-        residual_in_fp32,
-        False,
-    )
-def rms_norm_gated(
-    x: torch.Tensor,
-    g: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    activation: str = "swish",
-    residual: torch.Tensor | None = None,
-    prenorm: bool = False,
-    residual_in_fp32: bool = False,
-    eps: float = 1e-6,
-):
-    return LayerNormGatedFunction.apply(
-        x,
-        g,
-        weight,
-        bias,
-        activation,
-        residual,
-        eps,
-        prenorm,
-        residual_in_fp32,
-        True,
-    )
-def layer_norm_swish_gate_linear(
-    x: torch.Tensor,
-    g: torch.Tensor,
-    norm_weight: torch.Tensor,
-    norm_bias: torch.Tensor,
-    linear_weight: torch.Tensor,
-    linear_bias: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    prenorm: bool = False,
-    residual_in_fp32: bool = False,
-    eps: float = 1e-6,
-):
-    return LayerNormGatedLinearFunction.apply(
-        x,
-        g,
-        norm_weight,
-        norm_bias,
-        linear_weight,
-        linear_bias,
-        residual,
-        eps,
-        prenorm,
-        residual_in_fp32,
-        False,
-    )
-def rms_norm_swish_gate_linear(
-    x,
-    g: torch.Tensor,
-    norm_weight: torch.Tensor,
-    norm_bias: torch.Tensor,
-    linear_weight: torch.Tensor,
-    linear_bias: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    prenorm: bool = False,
-    residual_in_fp32: bool = False,
-    eps: float = 1e-6,
-):
-    return LayerNormGatedLinearFunction.apply(
-        x,
-        g,
-        norm_weight,
-        norm_bias,
-        linear_weight,
-        linear_bias,
-        residual,
-        eps,
-        prenorm,
-        residual_in_fp32,
-        True,
-    )
-class FusedLayerNormGated(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        bias: bool = False,
-        activation: str = "swish",
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormGated:
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.elementwise_affine = elementwise_affine
-        self.eps = eps
-        self.activation = activation
-        if self.activation not in ["swish", "silu", "sigmoid"]:
-            raise ValueError(f"Unsupported activation: {self.activation}")
-        self.register_parameter("weight", None)
-        self.register_parameter("bias", None)
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
-            if bias:
-                self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
-        self.reset_parameters()
-    def reset_parameters(self):
-        if self.elementwise_affine:
-            nn.init.ones_(self.weight)
-            if self.bias is not None:
-                nn.init.zeros_(self.bias)
-    def __repr__(self) -> str:
-        s = f"{self.__class__.__name__}({self.hidden_size}"
-        if not self.elementwise_affine:
-            s += f", elementwise_affine={self.elementwise_affine}"
-        s += f", eps={self.eps}"
-        s += f", activation={self.activation}"
-        s += ")"
-        return s
-    def forward(
-        self,
-        x: torch.Tensor,
-        g: torch.Tensor,
-        residual: torch.Tensor | None = None,
-        prenorm: bool = False,
-        residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
-        return layer_norm_gated(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            self.activation,
-            residual=residual,
-            eps=self.eps,
-            prenorm=prenorm,
-            residual_in_fp32=residual_in_fp32,
-        )
-class FusedRMSNormGated(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        activation: str = "swish",
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormGated:
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.elementwise_affine = elementwise_affine
-        self.eps = eps
-        self.activation = activation
-        if self.activation not in ["swish", "silu", "sigmoid"]:
-            raise ValueError(f"Unsupported activation: {self.activation}")
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
-        else:
-            self.register_parameter("weight", None)
-        self.register_parameter("bias", None)
-        self.reset_parameters()
-    def reset_parameters(self):
-        if self.elementwise_affine:
-            nn.init.ones_(self.weight)
-    def __repr__(self) -> str:
-        s = f"{self.__class__.__name__}({self.hidden_size}"
-        if not self.elementwise_affine:
-            s += f", elementwise_affine={self.elementwise_affine}"
-        s += f", eps={self.eps}"
-        s += f", activation={self.activation}"
-        s += ")"
-        return s
-    def forward(
-        self,
-        x: torch.Tensor,
-        g: torch.Tensor,
-        residual: torch.Tensor | None = None,
-        prenorm: bool = False,
-        residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
-        return rms_norm_gated(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            self.activation,
-            residual=residual,
-            eps=self.eps,
-            prenorm=prenorm,
-            residual_in_fp32=residual_in_fp32,
-        )
-class FusedLayerNormSwishGate(FusedLayerNormGated):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        bias: bool = False,
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormSwishGate:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            bias=bias,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )
-class FusedRMSNormSwishGate(FusedRMSNormGated):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormSwishGate:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )
-class FusedLayerNormGatedLinear(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormGatedLinear:
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.elementwise_affine = elementwise_affine
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
-        else:
-            self.register_parameter("weight", None)
-        self.register_parameter("bias", None)
-        self.reset_parameters()
-    def reset_parameters(self):
-        if self.elementwise_affine:
-            nn.init.ones_(self.weight)
-    def __repr__(self) -> str:
-        s = f"{self.__class__.__name__}({self.hidden_size}"
-        if not self.elementwise_affine:
-            s += f", elementwise_affine={self.elementwise_affine}"
-        s += f", eps={self.eps}"
-        s += ")"
-        return s
-    def forward(
-        self,
-        x: torch.Tensor,
-        g: torch.Tensor,
-        weight: torch.Tensor | None = None,
-        bias: torch.Tensor | None = None,
-        residual: torch.Tensor | None = None,
-        prenorm: bool = False,
-        residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
-        return layer_norm_swish_gate_linear(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            weight,
-            bias,
-            residual=residual,
-            eps=self.eps,
-            prenorm=prenorm,
-            residual_in_fp32=residual_in_fp32,
-        )
-class FusedLayerNormSwishGateLinear(FusedLayerNormGatedLinear):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormSwishGateLinear:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )
-class FusedRMSNormGatedLinear(nn.Module):
-    def __init__(
-        self,
-        hidden_size,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormGatedLinear:
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.elementwise_affine = elementwise_affine
-        self.eps = eps
-        self.register_parameter("weight", None)
-        self.register_parameter("bias", None)
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
-        self.reset_parameters()
-    def reset_parameters(self):
-        if self.elementwise_affine:
-            nn.init.ones_(self.weight)
-    def __repr__(self) -> str:
-        s = f"{self.__class__.__name__}({self.hidden_size}"
-        if not self.elementwise_affine:
-            s += f", elementwise_affine={self.elementwise_affine}"
-        s += f", eps={self.eps}"
-        s += ")"
-        return s
-    def forward(
-        self,
-        x: torch.Tensor,
-        g: torch.Tensor,
-        weight: torch.Tensor | None = None,
-        bias: torch.Tensor | None = None,
-        residual: torch.Tensor | None = None,
-        prenorm: bool = False,
-        residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
-        return rms_norm_swish_gate_linear(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            weight,
-            bias,
-            residual=residual,
-            eps=self.eps,
-            prenorm=prenorm,
-            residual_in_fp32=residual_in_fp32,
-        )
-class FusedRMSNormSwishGateLinear(FusedRMSNormGatedLinear):
-    def __init__(
-        self,
-        hidden_size: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormSwishGateLinear:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )

build/{torch-xpu/modules/conv/cp → torch-cuda/modules/grpo}/__init__.py RENAMED Viewed

@@ -5,9 +5,7 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-from .ops import CausalConv1dFunctionCP, causal_conv1d_cp
-__all__ = [
-    'CausalConv1dFunctionCP',
-    'causal_conv1d_cp',
-]

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+from ...modules.grpo.ops import fused_grpo_loss, grpo_loss_with_old_logps
+__all__ = ['fused_grpo_loss', 'grpo_loss_with_old_logps']

build/torch-cuda/modules/{grpo.py → grpo/ops.py} RENAMED Viewed

@@ -6,65 +6,20 @@
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
 # modified from https://github.com/mdy666/mdy_triton/blob/e0a856347bd988e05e0152332bba35f1d33c5b1f/others/grpo/grpo_loss.ipynb
-# XHS ID: blueeeee
-# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py
-"""
-# Get the per-token log probabilities for the completions for the model and the reference model
-    def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep):
-        # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
-        logits = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits
-        logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
-        input_ids = input_ids[:, -logits_to_keep:]
-        # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
-        # See https://github.com/huggingface/trl/issues/2770
-        logits = logits[:, -logits_to_keep:]
-        return selective_log_softmax(logits, input_ids)  #  compute logprobs for the input tokens
-    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
-        if return_outputs:
-            raise ValueError("The GRPOTrainer does not support returning outputs")
-        # Compute the per-token log probabilities for the model
-        prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
-        completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"]
-        input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
-        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
-        logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens
-        per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
-        # Compute the KL divergence between the model and the reference model
-        ref_per_token_logps = inputs["ref_per_token_logps"]
-        per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
-        # x - x.detach() allows for preserving gradients from x
-        advantages = inputs["advantages"]
-        per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
-        per_token_loss = -(per_token_loss - self.beta * per_token_kl)
-        loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
-        # Log the metrics
-        completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
-        self._metrics["completion_length"].append(completion_length)
-        mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
-        self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
-        return loss
-"""
 import torch
 import triton
 import triton.language as tl
-from ..modules.backends import dispatch
-from ..ops.utils.op import exp, log
-from ..utils import IS_AMD, autotune_cache_kwargs, input_guard
 NUM_WARPS_AUTOTUNE = [4, 8, 16] if IS_AMD else [4, 8, 16, 32]
 @triton.autotune(
@@ -72,7 +27,7 @@ NUM_WARPS_AUTOTUNE = [4, 8, 16] if IS_AMD else [4, 8, 16, 32]
         triton.Config({'BLOCK_SIZE': BLOCK_SIZE},  num_warps=NUM_WARPS, num_stages=NUM_STAGES)
         for BLOCK_SIZE in [1024, 2048, 4096, 8192]
         for NUM_WARPS in NUM_WARPS_AUTOTUNE
-        for NUM_STAGES in [1, 2, 4]
     ],
     key=['B', 'N'],
     **autotune_cache_kwargs,
@@ -198,8 +153,7 @@ def grpo_bwd_kernel(
         tl.debug_barrier()
         logp = x - lse
-        dlogp = (beta * (-1.0 * exp(ref_logp - logp) + 1)
-                 - advantage) * dloss
         for start_n in tl.range(0, N, BLOCK_SIZE):
             cols = start_n + base_cols
@@ -251,7 +205,10 @@ class GrpoLoss(torch.autograd.Function):
             lse_ptr=lse,
             beta=beta,
             save_kl=save_kl,
-            B=B, M=M, N=N, L=L,
             start_idx=input_ids_start_index,
         )
         ctx.beta = beta
@@ -263,7 +220,6 @@ class GrpoLoss(torch.autograd.Function):
     @input_guard
     @staticmethod
     def backward(ctx, dloss):
-        # The grad of logits comes from two parts, the reward part and the kl part
         lse, logits, input_ids, advantages, completion_mask = ctx.saved_tensors
         inplace = ctx.inplace
         B, L_ADD_1, N = ctx.input_shape
@@ -286,46 +242,52 @@ class GrpoLoss(torch.autograd.Function):
             completion_mask_ptr=completion_mask,
             lse_ptr=lse,
             beta=ctx.beta,
-            B=B, N=N, L=L,
-            BLOCK_SIZE=BN,
             start_idx=input_ids_start_index,
         )
-        # The last token in the completion is not used in the loss computation
-        # and therefore its gradient should be set to 0
         dlogits[:, -1, :].fill_(0.0)
         return dlogits.view(*ctx.input_shape), None, None, None, None, None, None, None
-@dispatch('modules')
-def fused_grpo_loss(logits, ref_logp, input_ids, advantages,
-                    beta=0.1, completion_mask=None, save_kl=False, inplace=False) -> torch.Tensor:
-    '''
-    compute grpo loss, save memory(no addition usage) and fast speed(6X for A800)
     Args:
-        logtits: Tensor, [B, L+1, vocab_size], the origin output of model, it's not logits[:, :-1]
-        ref_logp: Tensor, [B, L], the origin output of model, it's not ref_logits[:, :-1]
-        input_ids: Tensor, [B, K+L], it's prompt_completion_id, it contains the prompt ids and output ids
-        advantages: Tensor, [B], the advantages of each prompt
-        beta: float, the weight of kl loss
-        completion_mask: Tensor, loss mask
-        save_kl: bool, if true will save kl
-    Retutn:
-        loss: Tensor, [B, L], the loss of grpo, it contains the advantage part and kl part
-    NOTE: logits(ref_logits) is computed by these steps
-        logits_to_keep = completion_ids.size(1)
-        def get_per_token_logits(model, input_ids, attention_mask, logits_to_keep):
-            # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
-            logits = model(
-                input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1
-            ).logits
-            return logits
-        logits = get_per_token_logits(model, prompt_completion_ids, attention_mask, logits_to_keep)
-    '''
     out = GrpoLoss.apply(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl, inplace)
     if not save_kl:
         return out
@@ -356,7 +318,8 @@ def grpo_loss_torch(logits, ref_logp, input_ids, advantages, beta=0.1, completio
     return per_token_loss if not save_kl else (per_token_loss, per_token_kl)
-@torch.compile(fullgraph=True)
 def grpo_loss_with_old_logps(
     logps: torch.Tensor,
     ref_logps: torch.Tensor,
@@ -367,55 +330,52 @@ def grpo_loss_with_old_logps(
     beta: float = 0.2,
     epsilon: float = 0.2,
 ):
-    """
-    Compute the GRPO (Group Relative Policy Optimization) loss.
     Args:
-        logps (torch.Tensor): [Batch, Token_length] Log probabilities of the current policy.
-        ref_logps (torch.Tensor):[Batch, Token_length]  Log probabilities of the reference policy.
-        old_logps (torch.Tensor): [Batch, Token_length] Log probabilities of the old policy.
-        completion_ids (torch.Tensor): [Batch, Token_length] Completion token IDs (bool).
-        pad_token_id: Pad token ID.
-        logits_to_keep (int): Number of logits to keep for masking.
-        rewards (torch.Tensor): [Batch] Rewards for each generation.
-        beta (float) = 0.2: A hyperparameter for weighting the KL divergence term.
-        epsilon (float) = 0.2: An float hyperparameter for clipping the importance weights.
     Returns:
-        torch.Tensor: The computed GRPO loss.
     """
     B = logps.shape[0]
     assert B > 1, "Batch * Num generations should be greater than 1"
     rewards_shaped = rewards.view(-1, B)  # B,num_generations
-    advantages = (rewards_shaped - rewards_shaped.mean(dim=1, keepdim=True)) / \
-        (rewards_shaped.std(dim=1, keepdim=True) + 1e-8)
     advantages = advantages.view(-1)  # B*num_generations
-    # Calculate the per - token KL divergence
     per_token_kl = torch.exp(ref_logps - logps) - (ref_logps - logps) - 1
-    # Calculate the ratio of probabilities (importance weights)
-    # Importance weights are calculated as exp(log_pi_theta - log_pi_theta_old)
     importance_weights = torch.exp(logps - old_logps)
-    # Clip the importance weights to the range [1 - epsilon, 1 + epsilon]
     importance_weights_clipped = torch.clamp(importance_weights, 1 - epsilon, 1 + epsilon)
-    # Create a completion mask. It checks which positions are valid based on logits_to_keep
     completion_mask = torch.arange(logits_to_keep, device=logps.device)[None, :] >= 0
-    # Combine the completion mask and padding mask
-    completion_mask = completion_mask & pad_mask  # Ensure matching shape
-    # Add an extra dimension to advantages to match the shape for element - wise multiplication
     advantages = advantages.unsqueeze(1)
-    # Calculate the per - token loss. It takes the minimum of the unclipped and clipped importance weights
-    # and subtracts the KL divergence term weighted by beta, then multiplies by the completion mask
-    token_loss = -(torch.min(advantages * importance_weights, advantages *
-                   importance_weights_clipped) - beta * per_token_kl) * completion_mask
-    # Calculate the final loss by summing the token losses and normalizing by the number of valid tokens
-    loss = -token_loss.sum() / completion_mask.sum()
     return loss

 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
 # modified from https://github.com/mdy666/mdy_triton/blob/e0a856347bd988e05e0152332bba35f1d33c5b1f/others/grpo/grpo_loss.ipynb
+# loss reference: https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py
 import torch
 import triton
 import triton.language as tl
+from ...backends import dispatch
+from ...ops.utils.op import exp, log
+from ...utils import IS_AMD, IS_INTEL, IS_NPU, autotune_cache_kwargs, input_guard
 NUM_WARPS_AUTOTUNE = [4, 8, 16] if IS_AMD else [4, 8, 16, 32]
+# single-stage pipelining avoids Intel/XPU correctness failures with cold memory.
+NUM_STAGES_AUTOTUNE = [1] if IS_INTEL else [1, 2, 4]
 @triton.autotune(
         triton.Config({'BLOCK_SIZE': BLOCK_SIZE},  num_warps=NUM_WARPS, num_stages=NUM_STAGES)
         for BLOCK_SIZE in [1024, 2048, 4096, 8192]
         for NUM_WARPS in NUM_WARPS_AUTOTUNE
+        for NUM_STAGES in NUM_STAGES_AUTOTUNE
     ],
     key=['B', 'N'],
     **autotune_cache_kwargs,
         tl.debug_barrier()
         logp = x - lse
+        dlogp = (beta * (-1.0 * exp(ref_logp - logp) + 1) - advantage) * dloss
         for start_n in tl.range(0, N, BLOCK_SIZE):
             cols = start_n + base_cols
             lse_ptr=lse,
             beta=beta,
             save_kl=save_kl,
+            B=B,
+            M=M,
+            N=N,
+            L=L,
             start_idx=input_ids_start_index,
         )
         ctx.beta = beta
     @input_guard
     @staticmethod
     def backward(ctx, dloss):
         lse, logits, input_ids, advantages, completion_mask = ctx.saved_tensors
         inplace = ctx.inplace
         B, L_ADD_1, N = ctx.input_shape
             completion_mask_ptr=completion_mask,
             lse_ptr=lse,
             beta=ctx.beta,
+            B=B,
+            N=N,
+            L=L,
             start_idx=input_ids_start_index,
+            BLOCK_SIZE=BN,
         )
+        # the final token has no loss contribution, so its gradient is zero.
         dlogits[:, -1, :].fill_(0.0)
         return dlogits.view(*ctx.input_shape), None, None, None, None, None, None, None
+@dispatch
+def fused_grpo_loss(
+    logits,
+    ref_logp,
+    input_ids,
+    advantages,
+    beta=0.1,
+    completion_mask=None,
+    save_kl=False,
+    inplace=False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Compute per-token GRPO loss from completion logits.
     Args:
+        logits (torch.Tensor):
+            Model logits of shape `[B, L + 1, V]`; the final token is excluded from the loss.
+        ref_logp (torch.Tensor):
+            Reference model log probabilities of shape `[B, L]` for the sampled completion tokens.
+        input_ids (torch.Tensor):
+            Prompt and completion token IDs of shape `[B, K + L]`.
+        advantages (torch.Tensor):
+            Per-sequence advantages of shape `[B]`.
+        beta (float, Optional):
+            KL penalty coefficient. Default: 0.1.
+        completion_mask (torch.Tensor, Optional):
+            Mask of shape `[B, L]`; masked tokens have zero loss and gradient. Default: `None`.
+        save_kl (bool, Optional):
+            Whether to also return the per-token KL penalty. Default: `False`.
+        inplace (bool, Optional):
+            Whether backward may overwrite logits with gradients. Default: `False`.
+    Returns:
+        torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+            Per-token loss of shape `[B, L]`, with a separate KL tensor when `save_kl=True`.
+    """
     out = GrpoLoss.apply(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl, inplace)
     if not save_kl:
         return out
     return per_token_loss if not save_kl else (per_token_loss, per_token_kl)
+# keep this path eager on NPU to avoid an Inductor miscompile.
+@torch.compile(fullgraph=True, disable=IS_NPU)
 def grpo_loss_with_old_logps(
     logps: torch.Tensor,
     ref_logps: torch.Tensor,
     beta: float = 0.2,
     epsilon: float = 0.2,
 ):
+    """Compute clipped GRPO loss using current, reference and old-policy log probabilities.
     Args:
+        logps (torch.Tensor):
+            Current policy log probabilities of shape `[B, T]`.
+        ref_logps (torch.Tensor):
+            Reference policy log probabilities of shape `[B, T]`.
+        old_logps (torch.Tensor):
+            Old policy log probabilities of shape `[B, T]`.
+        pad_mask (torch.Tensor):
+            Boolean mask selecting completion tokens of shape `[B, T]`.
+        logits_to_keep (int):
+            Number of completion tokens, equal to `T`.
+        rewards (torch.Tensor):
+            Per-sequence rewards used to normalize advantages.
+        beta (float, Optional):
+            KL penalty coefficient. Default: 0.2.
+        epsilon (float, Optional):
+            Clipping range around an importance weight of one. Default: 0.2.
     Returns:
+        torch.Tensor:
+            Scalar loss averaged over valid completion tokens.
     """
     B = logps.shape[0]
     assert B > 1, "Batch * Num generations should be greater than 1"
     rewards_shaped = rewards.view(-1, B)  # B,num_generations
+    advantages = (rewards_shaped - rewards_shaped.mean(dim=1, keepdim=True)) / (rewards_shaped.std(dim=1, keepdim=True) + 1e-8)
     advantages = advantages.view(-1)  # B*num_generations
     per_token_kl = torch.exp(ref_logps - logps) - (ref_logps - logps) - 1
     importance_weights = torch.exp(logps - old_logps)
     importance_weights_clipped = torch.clamp(importance_weights, 1 - epsilon, 1 + epsilon)
     completion_mask = torch.arange(logits_to_keep, device=logps.device)[None, :] >= 0
+    completion_mask = completion_mask & pad_mask
     advantages = advantages.unsqueeze(1)
+    token_loss = -(
+        torch.min(advantages * importance_weights, advantages * importance_weights_clipped) - beta * per_token_kl
+    ) * completion_mask
+    loss = token_loss.sum() / completion_mask.sum()
     return loss

build/torch-cuda/modules/l2warp.py CHANGED Viewed

@@ -39,13 +39,12 @@ class L2Wrap(torch.autograd.Function):
         return loss
     @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):
         maxx, ids = ctx.saved_tensors
-        glogits = torch.zeros(ctx.logits_shape, device=grad_output.device, dtype=grad_output.dtype)
-        # an autograd.Function must scale its input gradients by the upstream gradient; fold the
-        # scalar grad_output into the sparse maxx to avoid a second full-size logits allocation
-        glogits.scatter_(-1, ids, maxx * grad_output)
-        return grad_output, glogits, None
 l2_warp = L2Wrap.apply

         return loss
     @staticmethod
+    def backward(ctx, dloss: torch.Tensor):
         maxx, ids = ctx.saved_tensors
+        dlogits = torch.zeros(ctx.logits_shape, device=dloss.device, dtype=maxx.dtype)
+        # scale before scattering to avoid another full-size logits allocation
+        dlogits.scatter_(-1, ids, (maxx * dloss).to(maxx.dtype))
+        return dloss, dlogits, None
 l2_warp = L2Wrap.apply

build/torch-cuda/modules/{conv/long_conv.py → long_conv.py} RENAMED Viewed

@@ -13,7 +13,13 @@ import torch.nn.functional as F
 from einops import rearrange
-def fft_conv(u, k, dropout_mask, gelu=True, k_rev=None):
     seqlen = u.shape[-1]
     fft_size = 2 * seqlen
     k_f = torch.fft.rfft(k, n=fft_size) / fft_size
@@ -37,60 +43,52 @@ def fft_conv(u, k, dropout_mask, gelu=True, k_rev=None):
 class LongConvolution(nn.Module):
     """
-    LongConvolution applies a convolution operation on the input tensor using a fixed
-    filter of length max_len.
-    The filter is learned during training and is applied using FFT convolution.
     Args:
-        hidden_size (int): The number of expected features in the input and output.
-        max_len (int): The maximum sequence length.
     Returns:
-        y: [batch_size, seq_len, hidden_size] tensor
     """
-    def __init__(
-        self,
-        hidden_size: int,
-        max_len: int,
-        **kwargs,
-    ):
-        """
-        Initializes the LongConvolution module.
-        Args:
-            hidden_size (int): The number of expected features in the input and output.
-            max_len (int): The maximum sequence length.
-        """
         super().__init__()
         self.hidden_size = hidden_size
         self.filter = nn.Parameter(torch.randn(self.hidden_size, max_len), requires_grad=True)
-    def forward(self, x: torch.Tensor, *args, **kwargs):
         """
-        Applies the LongConvolution operation on the input tensor.
         Args:
-            x: [batch_size, seq_len, hidden_size] tensor
         Returns:
-            y: [batch_size, seq_len, hidden_size] tensor
         """
         x = x.transpose(1, 2)
-        y = fft_conv(x, self.filter, dropout_mask=None, gelu=False)
         y = y.transpose(1, 2)
         return y.to(dtype=x.dtype)
 class PositionalEmbedding(nn.Module):
-    def __init__(self, emb_dim: int, seq_len: int, **kwargs):
         """Complex exponential positional embeddings for implicit long convolution filters."""
         super().__init__()
         self.seq_len = seq_len
-        # The time embedding fed to the filteres is normalized so that t_f = 1
         t = torch.linspace(0, 1, self.seq_len)[None, :, None]  # 1, L, 1
         if emb_dim > 1:
             bands = (emb_dim - 1) // 2
-        # To compute the right embeddings we use the "proper" linspace
         t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]
         w = 2 * math.pi * t_rescaled / seq_len  # 1, L, 1
@@ -99,7 +97,7 @@ class PositionalEmbedding(nn.Module):
         z = torch.cat([t, z.real, z.imag], dim=-1)
         self.z = nn.Parameter(z, requires_grad=False)
-    def forward(self, L):
         return self.z[:, :L]
@@ -112,61 +110,45 @@ class ImplicitLongConvolution(nn.Module):
             The number of expected features in the input and output.
         max_len (int):
             The maximum sequence length.
-        d_emb (Optional[int]):
             The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine).
-            Defaults to 3.
-        d_hidden (Optional[int]):
-            The number of features in the hidden layer of the MLP. Defaults to 16.
     Attributes:
-        pos_emb (`PositionalEmbedding`): The positional embedding layer.
-        mlp (`nn.Sequential`): The MLP that parameterizes the implicit filter.
     """
-    def __init__(
-        self,
-        hidden_size: int,
-        max_len: int,
-        d_emb: int = 3,
-        d_hidden: int = 16,
-        **kwargs,
-    ):
-        """
-        Long convolution with implicit filter parameterized by an MLP.
-        """
         super().__init__()
         self.hidden_size = hidden_size
         self.d_emb = d_emb
-        assert (
-            d_emb % 2 != 0 and d_emb >= 3
-        ), "d_emb must be odd and greater or equal to 3 (time, sine and cosine)"
-        self.pos_emb = PositionalEmbedding(d_emb, max_len)
-        # final linear layer
-        self.mlp = nn.Sequential(
-            nn.Linear(d_emb, d_hidden),
-            torch.nn.ReLU(),
-            nn.Linear(d_hidden, hidden_size),
-        )
-    def filter(self, seq_len: int, *args, **kwargs):
         return self.mlp(self.pos_emb(seq_len)).transpose(1, 2)
-    def forward(self, x: torch.Tensor, *args, **kwargs):
         """
         Args:
-            x: [batch_size, seq_len, hidden_size] tensor
         Returns:
-            y: [batch_size, seq_len, hidden_size] tensor
         """
         x = x.transpose(1, 2)
         k = self.filter(x.shape[-1])
-        y = fft_conv(x, k, dropout_mask=None, gelu=False)
         y = y.transpose(1, 2)
         return y.to(dtype=x.dtype)

 from einops import rearrange
+def fft_conv(
+    u: torch.Tensor,
+    k: torch.Tensor,
+    dropout_mask: torch.Tensor | None,
+    gelu: bool = True,
+    k_rev: torch.Tensor | None = None,
+) -> torch.Tensor:
     seqlen = u.shape[-1]
     fft_size = 2 * seqlen
     k_f = torch.fft.rfft(k, n=fft_size) / fft_size
 class LongConvolution(nn.Module):
     """
+    Apply FFT convolution with a learned filter of length max_len.
     Args:
+        hidden_size (int):
+            The number of expected features in the input and output.
+        max_len (int):
+            The maximum sequence length.
     Returns:
+        y (torch.Tensor):
+            Output of shape [batch_size, seq_len, hidden_size].
     """
+    def __init__(self, hidden_size: int, max_len: int, **kwargs) -> None:
         super().__init__()
         self.hidden_size = hidden_size
         self.filter = nn.Parameter(torch.randn(self.hidden_size, max_len), requires_grad=True)
+    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         """
         Args:
+            x (torch.Tensor):
+                Input of shape [batch_size, seq_len, hidden_size].
         Returns:
+            y (torch.Tensor):
+                Output of shape [batch_size, seq_len, hidden_size].
         """
         x = x.transpose(1, 2)
+        y = fft_conv(u=x, k=self.filter, dropout_mask=None, gelu=False)
         y = y.transpose(1, 2)
         return y.to(dtype=x.dtype)
 class PositionalEmbedding(nn.Module):
+    def __init__(self, emb_dim: int, seq_len: int, **kwargs) -> None:
         """Complex exponential positional embeddings for implicit long convolution filters."""
         super().__init__()
         self.seq_len = seq_len
+        # the time embedding fed to the filters is normalized so that t_f = 1
         t = torch.linspace(0, 1, self.seq_len)[None, :, None]  # 1, L, 1
         if emb_dim > 1:
             bands = (emb_dim - 1) // 2
+        # use integer time steps for the complex exponential embeddings
         t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]
         w = 2 * math.pi * t_rescaled / seq_len  # 1, L, 1
         z = torch.cat([t, z.real, z.imag], dim=-1)
         self.z = nn.Parameter(z, requires_grad=False)
+    def forward(self, L: int) -> torch.Tensor:
         return self.z[:, :L]
             The number of expected features in the input and output.
         max_len (int):
             The maximum sequence length.
+        d_emb (int, Optional):
             The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine).
+            Default: 3.
+        d_hidden (int, Optional):
+            The number of features in the hidden layer of the MLP. Default: 16.
     Attributes:
+        pos_emb (PositionalEmbedding):
+            The positional embedding layer.
+        mlp (nn.Sequential):
+            The MLP that parameterizes the implicit filter.
     """
+    def __init__(self, hidden_size: int, max_len: int, d_emb: int = 3, d_hidden: int = 16, **kwargs) -> None:
         super().__init__()
         self.hidden_size = hidden_size
         self.d_emb = d_emb
+        assert d_emb % 2 != 0 and d_emb >= 3, "d_emb must be odd and greater or equal to 3 (time, sine and cosine)"
+        self.pos_emb = PositionalEmbedding(emb_dim=d_emb, seq_len=max_len)
+        self.mlp = nn.Sequential(nn.Linear(d_emb, d_hidden), torch.nn.ReLU(), nn.Linear(d_hidden, hidden_size))
+    def filter(self, seq_len: int, *args, **kwargs) -> torch.Tensor:
         return self.mlp(self.pos_emb(seq_len)).transpose(1, 2)
+    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         """
         Args:
+            x (torch.Tensor):
+                Input of shape [batch_size, seq_len, hidden_size].
         Returns:
+            y (torch.Tensor):
+                Output of shape [batch_size, seq_len, hidden_size].
         """
         x = x.transpose(1, 2)
         k = self.filter(x.shape[-1])
+        y = fft_conv(u=x, k=k, dropout_mask=None, gelu=False)
         y = y.transpose(1, 2)
         return y.to(dtype=x.dtype)

build/{torch-xpu/modules → torch-cuda/modules/norm}/__init__.py RENAMED Viewed

@@ -5,48 +5,46 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
-from ..modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution
-from ..modules.fused_bitlinear import BitLinear, FusedBitLinear
-from ..modules.fused_cross_entropy import FusedCrossEntropyLoss
-from ..modules.fused_kl_div import FusedKLDivLoss
-from ..modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
-from ..modules.fused_norm_gate import (
     FusedLayerNormGated,
     FusedLayerNormSwishGate,
     FusedLayerNormSwishGateLinear,
     FusedRMSNormGated,
     FusedRMSNormSwishGate,
     FusedRMSNormSwishGateLinear,
 )
-from ..modules.l2norm import L2Norm
-from ..modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear
-from ..modules.mlp import GatedMLP
-from ..modules.rotary import RotaryEmbedding
-from ..modules.token_shift import TokenShift
 __all__ = [
-    'BitLinear',
-    'FusedBitLinear',
-    'FusedCrossEntropyLoss',
-    'FusedKLDivLoss',
     'FusedLayerNormGated',
     'FusedLayerNormSwishGate',
     'FusedLayerNormSwishGateLinear',
-    'FusedLinearCrossEntropyLoss',
     'FusedRMSNormGated',
     'FusedRMSNormSwishGate',
     'FusedRMSNormSwishGateLinear',
-    'GatedMLP',
     'GroupNorm',
     'GroupNormLinear',
-    'ImplicitLongConvolution',
     'L2Norm',
     'LayerNorm',
     'LayerNormLinear',
-    'LongConvolution',
     'RMSNorm',
     'RMSNormLinear',
-    'RotaryEmbedding',
-    'ShortConvolution',
-    'TokenShift',
 ]

 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+from ...modules.norm.fused_norm_gate import (
     FusedLayerNormGated,
+    FusedLayerNormGatedLinear,
     FusedLayerNormSwishGate,
     FusedLayerNormSwishGateLinear,
     FusedRMSNormGated,
+    FusedRMSNormGatedLinear,
     FusedRMSNormSwishGate,
     FusedRMSNormSwishGateLinear,
 )
+from ...modules.norm.l2norm import L2Norm
+from ...modules.norm.layernorm import (
+    GroupNorm,
+    GroupNormLinear,
+    LayerNorm,
+    LayerNormLinear,
+    NormParallel,
+    RMSNorm,
+    RMSNormLinear,
+)
+from ...modules.norm.layernorm_gated import LayerNormGated, RMSNormGated
 __all__ = [
     'FusedLayerNormGated',
+    'FusedLayerNormGatedLinear',
     'FusedLayerNormSwishGate',
     'FusedLayerNormSwishGateLinear',
     'FusedRMSNormGated',
+    'FusedRMSNormGatedLinear',
     'FusedRMSNormSwishGate',
     'FusedRMSNormSwishGateLinear',
     'GroupNorm',
     'GroupNormLinear',
     'L2Norm',
     'LayerNorm',
+    'LayerNormGated',
     'LayerNormLinear',
+    'NormParallel',
     'RMSNorm',
+    'RMSNormGated',
     'RMSNormLinear',
 ]

build/{torch-xpu/modules → torch-cuda/modules/norm}/fused_norm_gate.py RENAMED Viewed

@@ -15,7 +15,8 @@ import torch.nn.functional as F
 import triton
 import triton.language as tl
-from ..utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard
 @triton.heuristics(
@@ -31,20 +32,20 @@ from ..utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard
     key=["D", "NB", "IS_RMS_NORM", "STORE_RESIDUAL_OUT", "HAS_RESIDUAL", "HAS_WEIGHT"],
     **autotune_cache_kwargs,
 )
-@triton.jit
 def layer_norm_gated_fwd_kernel(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    y,  # pointer to the output
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    residual,  # pointer to the residual
-    residual_out,  # pointer to the residual
-    mean,  # pointer to the mean
-    rstd,  # pointer to the 1/std
-    eps,  # epsilon to avoid division by zero
-    T,  # number of rows in x
-    D: tl.constexpr,  # number of columns in x
     BT: tl.constexpr,
     BD: tl.constexpr,
     NB: tl.constexpr,
@@ -55,23 +56,26 @@ def layer_norm_gated_fwd_kernel(
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
 ):
-    i_t = tl.program_id(0)
     o_d = tl.arange(0, BD)
     m_d = o_d < D
-    p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-    b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
     if HAS_RESIDUAL:
-        p_res = tl.make_block_ptr(residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-        b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32)
     if STORE_RESIDUAL_OUT:
-        p_res_out = tl.make_block_ptr(residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-        tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1))
     if not IS_RMS_NORM:
         b_mean = tl.sum(b_x, axis=1) / D
-        p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,))
-        tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,))
         b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)
         b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
     else:
@@ -79,8 +83,8 @@ def layer_norm_gated_fwd_kernel(
         b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
     b_rstd = 1 / tl.sqrt(b_var + eps)
-    p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))
-    tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))
     if HAS_WEIGHT:
         b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
@@ -92,16 +96,15 @@ def layer_norm_gated_fwd_kernel(
         b_y = b_y + b_b[None, :]
     # swish/sigmoid output gate
-    p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-    b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
     if ACTIVATION == "swish" or ACTIVATION == "silu":
         b_y = b_y * b_g * tl.sigmoid(b_g)
     elif ACTIVATION == "sigmoid":
         b_y = b_y * tl.sigmoid(b_g)
-    # Write output
-    p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
-    tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
 @triton.heuristics(
@@ -118,18 +121,18 @@ def layer_norm_gated_fwd_kernel(
     **autotune_cache_kwargs,
 )
 @triton.jit
-def layer_norm_gated_fwd_kernel1(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    y,  # pointer to the output
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    residual,  # pointer to the residual
-    residual_out,  # pointer to the residual
-    mean,  # pointer to the mean
-    rstd,  # pointer to the 1/std
-    eps,  # epsilon to avoid division by zero
-    D: tl.constexpr,  # number of columns in x
     BD: tl.constexpr,
     ACTIVATION: tl.constexpr,
     IS_RMS_NORM: tl.constexpr,
@@ -138,7 +141,7 @@ def layer_norm_gated_fwd_kernel1(
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
 ):
-    i_t = tl.program_id(0)
     x += i_t * D
     y += i_t * D
     g += i_t * D
@@ -181,7 +184,6 @@ def layer_norm_gated_fwd_kernel1(
     elif ACTIVATION == "sigmoid":
         b_y = b_y * tl.sigmoid(b_g)
-    # Write output
     tl.store(y + o_d, b_y, mask=m_d)
@@ -198,18 +200,18 @@ def layer_norm_gated_fwd_kernel1(
     key=["D", "NB", "IS_RMS_NORM", "HAS_DRESIDUAL", "HAS_WEIGHT"],
     **autotune_cache_kwargs,
 )
-@triton.jit
 def layer_norm_gated_bwd_kernel(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    y,  # pointer to the output to be recomputed
-    dy,  # pointer to the output gradient
-    dx,  # pointer to the input gradient
-    dg,  # pointer to the gate gradient
-    dw,  # pointer to the partial sum of weights gradient
-    db,  # pointer to the partial sum of biases gradient
     dresidual,
     dresidual_in,
     mean,
@@ -228,7 +230,7 @@ def layer_norm_gated_bwd_kernel(
     HAS_BIAS: tl.constexpr,
     RECOMPUTE_OUTPUT: tl.constexpr,
 ):
-    i_s = tl.program_id(0)
     o_d = tl.arange(0, BD)
     m_d = o_d < D
     if HAS_WEIGHT:
@@ -238,45 +240,45 @@ def layer_norm_gated_bwd_kernel(
         b_b = tl.load(b + o_d, mask=m_d, other=0.0).to(tl.float32)
         b_db = tl.zeros((BT, BD), dtype=tl.float32)
-    # the caller guarantees NS = min(SM, T), so every program has at least one token.
-    # the last program's range may slightly exceed T (since BS = ceil(T/NS));
-    # make_block_ptr uses the true tensor shape (T, D), so boundary_check
-    # handles the partial tail tile by zero-padding loads and skipping stores.
-    # the m_t mask below further ensures dw/db only accumulate valid rows (< T).
     for i_t in range(i_s * BS, i_s * BS + BS, BT):
-        p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-        p_dg = tl.make_block_ptr(dg, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
         # [BT, BD]
-        b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
-        b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
-        b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32)
         if not IS_RMS_NORM:
-            p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t,), (BT,), (0,))
-            b_mean = tl.load(p_mean, boundary_check=(0,))
-        p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t,), (BT,), (0,))
-        b_rstd = tl.load(p_rstd, boundary_check=(0,))
-        # Compute dx
         b_xhat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None]
         b_xhat = tl.where(m_d[None, :], b_xhat, 0.0)
         b_y = b_xhat * b_w[None, :] if HAS_WEIGHT else b_xhat
         if HAS_BIAS:
             b_y = b_y + b_b[None, :]
-        if RECOMPUTE_OUTPUT:
-            p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-            tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
         b_sigmoid_g = tl.sigmoid(b_g)
         if ACTIVATION == "swish" or ACTIVATION == "silu":
             b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
-            b_dy = b_dy * b_g * b_sigmoid_g
         elif ACTIVATION == "sigmoid":
             b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
-            b_dy = b_dy * b_sigmoid_g
         b_wdy = b_dy
         if HAS_WEIGHT or HAS_BIAS:
@@ -296,16 +298,16 @@ def layer_norm_gated_bwd_kernel(
             b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D
             b_dx = (b_wdy - b_xhat * b_c1[:, None]) * b_rstd[:, None]
         if HAS_DRESIDUAL:
-            p_dres = tl.make_block_ptr(dresidual, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-            b_dres = tl.load(p_dres, boundary_check=(0, 1)).to(tl.float32)
             b_dx += b_dres
-        # Write dx
         if STORE_DRESIDUAL:
-            p_dres_in = tl.make_block_ptr(dresidual_in, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0))
-            tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1))
     if HAS_WEIGHT:
         tl.store(dw + i_s * D + o_d, tl.sum(b_dw, axis=0), mask=m_d)
@@ -326,18 +328,18 @@ def layer_norm_gated_bwd_kernel(
     key=["D", "IS_RMS_NORM", "STORE_DRESIDUAL", "HAS_DRESIDUAL", "HAS_WEIGHT"],
     **autotune_cache_kwargs,
 )
-@triton.jit
-def layer_norm_gated_bwd_kernel1(
-    x,  # pointer to the input
-    g,  # pointer to the gate
-    w,  # pointer to the weights
-    b,  # pointer to the biases
-    y,  # pointer to the output to be recomputed
-    dy,  # pointer to the output gradient
-    dx,  # pointer to the input gradient
-    dg,  # pointer to the gate gradient
-    dw,  # pointer to the partial sum of weights gradient
-    db,  # pointer to the partial sum of biases gradient
     dresidual,
     dresidual_in,
     mean,
@@ -354,7 +356,7 @@ def layer_norm_gated_bwd_kernel1(
     HAS_BIAS: tl.constexpr,
     RECOMPUTE_OUTPUT: tl.constexpr,
 ):
-    i_s = tl.program_id(0)
     o_d = tl.arange(0, BD)
     mask = o_d < D
     x += i_s * BS * D
@@ -376,7 +378,7 @@ def layer_norm_gated_bwd_kernel1(
         b_db = tl.zeros((BD,), dtype=tl.float32)
     for i_t in range(i_s * BS, min(i_s * BS + BS, T)):
-        # Load data to SRAM
         b_x = tl.load(x + o_d, mask=mask, other=0).to(tl.float32)
         b_g = tl.load(g + o_d, mask=mask, other=0).to(tl.float32)
         b_dy = tl.load(dy + o_d, mask=mask, other=0).to(tl.float32)
@@ -384,23 +386,25 @@ def layer_norm_gated_bwd_kernel1(
         if not IS_RMS_NORM:
             b_mean = tl.load(mean + i_t)
         b_rstd = tl.load(rstd + i_t)
-        # Compute dx
         b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
         b_xhat = tl.where(mask, b_xhat, 0.0)
         b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat
         if HAS_BIAS:
             b_y = b_y + b_b
-        if RECOMPUTE_OUTPUT:
-            tl.store(y + o_d, b_y, mask=mask)
         b_sigmoid_g = tl.sigmoid(b_g)
         if ACTIVATION == "swish" or ACTIVATION == "silu":
             b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
-            b_dy = b_dy * b_g * b_sigmoid_g
         elif ACTIVATION == "sigmoid":
             b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
-            b_dy = b_dy * b_sigmoid_g
         b_wdy = b_dy
         if HAS_WEIGHT:
             b_wdy = b_dy * b_w
@@ -417,7 +421,7 @@ def layer_norm_gated_bwd_kernel1(
         if HAS_DRESIDUAL:
             b_dres = tl.load(dresidual + o_d, mask=mask, other=0).to(tl.float32)
             b_dx += b_dres
-        # Write dx
         if STORE_DRESIDUAL:
             tl.store(dresidual_in + o_d, b_dx, mask=mask)
         tl.store(dx + o_d, b_dx, mask=mask)
@@ -440,6 +444,8 @@ def layer_norm_gated_bwd_kernel1(
         tl.store(db + i_s * D + o_d, b_db, mask=mask)
 def layer_norm_gated_fwd(
     x: torch.Tensor,
     g: torch.Tensor,
@@ -447,9 +453,9 @@ def layer_norm_gated_fwd(
     bias: torch.Tensor,
     activation: str = "swish",
     eps: float = 1e-5,
-    residual: torch.Tensor = None,
-    out_dtype: torch.dtype = None,
-    residual_dtype: torch.dtype = None,
     is_rms_norm: bool = False,
 ):
     if residual is not None:
@@ -461,7 +467,7 @@ def layer_norm_gated_fwd(
         assert weight.shape == (D,)
     if bias is not None:
         assert bias.shape == (D,)
-    # allocate output
     y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
     if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype):
         residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
@@ -469,17 +475,13 @@ def layer_norm_gated_fwd(
         residual_out = None
     mean = torch.empty((T,), dtype=torch.float, device=x.device) if not is_rms_norm else None
     rstd = torch.empty((T,), dtype=torch.float, device=x.device)
-    # Less than 64KB per feature: enqueue fused kernel
     MAX_FUSED_SIZE = 65536 // x.element_size()
     BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
     if D > BD:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-    # heuristics for number of warps
     if D <= 512:
-        # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range
-        # of T before recompiling the kernel.
-        # NB = triton.cdiv(T, 2048)
         NB = triton.cdiv(T, 2048 * 32)
         def grid(meta):
@@ -504,7 +506,7 @@ def layer_norm_gated_fwd(
             IS_RMS_NORM=is_rms_norm,
         )
     else:
-        layer_norm_gated_fwd_kernel1[(T,)](
             x=x,
             g=g,
             y=y,
@@ -524,6 +526,8 @@ def layer_norm_gated_fwd(
     return y, mean, rstd, residual_out if residual_out is not None else x
 def layer_norm_gated_bwd(
     dy: torch.Tensor,
     x: torch.Tensor,
@@ -532,12 +536,12 @@ def layer_norm_gated_bwd(
     bias: torch.Tensor,
     activation: str = "swish",
     eps: float = 1e-5,
-    mean: torch.Tensor = None,
-    rstd: torch.Tensor = None,
-    dresidual: torch.Tensor = None,
     has_residual: bool = False,
     is_rms_norm: bool = False,
-    x_dtype: torch.dtype = None,
     recompute_output: bool = False,
 ):
     T, D = x.shape
@@ -548,20 +552,17 @@ def layer_norm_gated_bwd(
         assert weight.shape == (D,)
     if bias is not None:
         assert bias.shape == (D,)
-    # allocate output
     dx = torch.empty_like(x) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
     dg = torch.empty_like(g) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
     dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None
     y = torch.empty(T, D, dtype=dy.dtype, device=dy.device) if recompute_output else None
-    # Less than 64KB per feature: enqueue fused kernel
     MAX_FUSED_SIZE = 65536 // x.element_size()
     BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
     if D > BD:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-    # cap program count to T so no program is completely idle.
-    # without this, high-SM GPUs (e.g. B200, 160 SMs) with small T would
-    # launch idle programs whose make_block_ptr offsets exceed the tensor shape.
     NS = min(get_multiprocessor_count(x.device.index), T)
     BS = math.ceil(T / NS)
@@ -570,9 +571,7 @@ def layer_norm_gated_bwd(
     grid = (NS,)
     if D <= 512:
-        # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range
-        # of T before recompiling the kernel.
-        # NB = triton.cdiv(T, 2048)
         NB = triton.cdiv(T, 2048 * 32)
         layer_norm_gated_bwd_kernel[grid](
@@ -591,8 +590,8 @@ def layer_norm_gated_bwd(
             mean=mean,
             rstd=rstd,
             T=T,
-            D=D,
             BS=BS,
             BD=BD,
             NB=NB,
             ACTIVATION=activation,
@@ -600,7 +599,7 @@ def layer_norm_gated_bwd(
             STORE_DRESIDUAL=dresidual_in is not None,
         )
     else:
-        layer_norm_gated_bwd_kernel1[grid](
             x=x,
             g=g,
             w=weight,
@@ -616,8 +615,8 @@ def layer_norm_gated_bwd(
             mean=mean,
             rstd=rstd,
             T=T,
-            D=D,
             BS=BS,
             BD=BD,
             ACTIVATION=activation,
             IS_RMS_NORM=is_rms_norm,
@@ -625,7 +624,7 @@ def layer_norm_gated_bwd(
         )
     dw = dw.sum(0).to(weight.dtype) if weight is not None else None
     db = db.sum(0).to(bias.dtype) if bias is not None else None
-    # Don't need to compute dresidual_in separately in this case
     if has_residual and dx.dtype == x.dtype:
         dresidual_in = dx
     return (dx, dg, dw, db, dresidual_in) if not recompute_output else (dx, dg, dw, db, dresidual_in, y)
@@ -649,7 +648,7 @@ class LayerNormGatedFunction(torch.autograd.Function):
     ):
         x_shape_og = x.shape
         g_shape_og = g.shape
-        # reshape input data into 2D tensor
         x = x.reshape(-1, x.shape[-1])
         g = g.reshape(-1, g.shape[-1])
         if residual is not None:
@@ -739,7 +738,7 @@ class LayerNormGatedLinearFunction(torch.autograd.Function):
     ):
         x_shape_og = x.shape
         g_shape_og = g.shape
-        # reshape input data into 2D tensor
         x = x.reshape(-1, x.shape[-1])
         g = g.reshape(-1, g.shape[-1])
         if residual is not None:
@@ -761,7 +760,7 @@ class LayerNormGatedLinearFunction(torch.autograd.Function):
         linear_weight = linear_weight.to(dtype)
         linear_bias = linear_bias.to(dtype) if linear_bias is not None else None
         out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias)
-        # We don't store y, will be recomputed in the backward pass to save memory
         ctx.save_for_backward(residual_out, g, norm_weight, norm_bias, linear_weight, mean, rstd)
         ctx.x_shape_og = x_shape_og
         ctx.g_shape_og = g_shape_og
@@ -829,18 +828,7 @@ def layer_norm_gated(
     residual_in_fp32: bool = False,
     eps: float = 1e-6,
 ):
-    return LayerNormGatedFunction.apply(
-        x,
-        g,
-        weight,
-        bias,
-        activation,
-        residual,
-        eps,
-        prenorm,
-        residual_in_fp32,
-        False,
-    )
 def rms_norm_gated(
@@ -854,18 +842,7 @@ def rms_norm_gated(
     residual_in_fp32: bool = False,
     eps: float = 1e-6,
 ):
-    return LayerNormGatedFunction.apply(
-        x,
-        g,
-        weight,
-        bias,
-        activation,
-        residual,
-        eps,
-        prenorm,
-        residual_in_fp32,
-        True,
-    )
 def layer_norm_swish_gate_linear(
@@ -932,7 +909,7 @@ class FusedLayerNormGated(nn.Module):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormGated:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -975,17 +952,17 @@ class FusedLayerNormGated(nn.Module):
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
         return layer_norm_gated(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            self.activation,
             residual=residual,
-            eps=self.eps,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
         )
@@ -998,7 +975,7 @@ class FusedRMSNormGated(nn.Module):
         activation: str = "swish",
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormGated:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -1038,17 +1015,17 @@ class FusedRMSNormGated(nn.Module):
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
         return rms_norm_gated(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            self.activation,
             residual=residual,
-            eps=self.eps,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
         )
@@ -1061,7 +1038,7 @@ class FusedLayerNormSwishGate(FusedLayerNormGated):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormSwishGate:
         super().__init__(
             hidden_size=hidden_size,
             elementwise_affine=elementwise_affine,
@@ -1080,14 +1057,8 @@ class FusedRMSNormSwishGate(FusedRMSNormGated):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormSwishGate:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )
 class FusedLayerNormGatedLinear(nn.Module):
@@ -1098,7 +1069,7 @@ class FusedLayerNormGatedLinear(nn.Module):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormGatedLinear:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -1135,18 +1106,18 @@ class FusedLayerNormGatedLinear(nn.Module):
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
         return layer_norm_swish_gate_linear(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            weight,
-            bias,
             residual=residual,
-            eps=self.eps,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
         )
@@ -1158,25 +1129,19 @@ class FusedLayerNormSwishGateLinear(FusedLayerNormGatedLinear):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedLayerNormSwishGateLinear:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )
 class FusedRMSNormGatedLinear(nn.Module):
     def __init__(
         self,
-        hidden_size,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormGatedLinear:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -1212,18 +1177,18 @@ class FusedRMSNormGatedLinear(nn.Module):
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
-    ) -> torch.Tensor:
         return rms_norm_swish_gate_linear(
-            x,
-            g,
-            self.weight,
-            self.bias,
-            weight,
-            bias,
             residual=residual,
-            eps=self.eps,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
         )
@@ -1235,11 +1200,23 @@ class FusedRMSNormSwishGateLinear(FusedRMSNormGatedLinear):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
-    ) -> FusedRMSNormSwishGateLinear:
-        super().__init__(
-            hidden_size=hidden_size,
-            elementwise_affine=elementwise_affine,
-            eps=eps,
-            device=device,
-            dtype=dtype,
-        )

 import triton
 import triton.language as tl
+from ...backends import dispatch
+from ...utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard
 @triton.heuristics(
     key=["D", "NB", "IS_RMS_NORM", "STORE_RESIDUAL_OUT", "HAS_RESIDUAL", "HAS_WEIGHT"],
     **autotune_cache_kwargs,
 )
+@triton.jit(do_not_specialize=['T'])
 def layer_norm_gated_fwd_kernel(
+    x,
+    g,
+    y,
+    w,
+    b,
+    residual,
+    residual_out,
+    mean,
+    rstd,
+    eps,
+    T,
+    D: tl.constexpr,
     BT: tl.constexpr,
     BD: tl.constexpr,
     NB: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
 ):
+    i_t = tl.program_id(0).to(tl.int64)
+    o_t = i_t * BT + tl.arange(0, BT)
     o_d = tl.arange(0, BD)
     m_d = o_d < D
+    m_t = o_t < T
+    m_x = m_t[:, None] & m_d[None, :]
+    p_x = x + o_t[:, None] * D + o_d[None, :]
+    b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32)
     if HAS_RESIDUAL:
+        p_res = residual + o_t[:, None] * D + o_d[None, :]
+        b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32)
     if STORE_RESIDUAL_OUT:
+        p_res_out = residual_out + o_t[:, None] * D + o_d[None, :]
+        tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x)
     if not IS_RMS_NORM:
         b_mean = tl.sum(b_x, axis=1) / D
+        p_mean = mean + o_t
+        tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t)
         b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)
         b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
     else:
         b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
     b_rstd = 1 / tl.sqrt(b_var + eps)
+    p_rstd = rstd + o_t
+    tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t)
     if HAS_WEIGHT:
         b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
         b_y = b_y + b_b[None, :]
     # swish/sigmoid output gate
+    p_g = g + o_t[:, None] * D + o_d[None, :]
+    b_g = tl.load(p_g, mask=m_x, other=0.0).to(tl.float32)
     if ACTIVATION == "swish" or ACTIVATION == "silu":
         b_y = b_y * b_g * tl.sigmoid(b_g)
     elif ACTIVATION == "sigmoid":
         b_y = b_y * tl.sigmoid(b_g)
+    p_y = y + o_t[:, None] * D + o_d[None, :]
+    tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x)
 @triton.heuristics(
     **autotune_cache_kwargs,
 )
 @triton.jit
+def layer_norm_gated_fwd_kernel_row(
+    x,
+    g,
+    y,
+    w,
+    b,
+    residual,
+    residual_out,
+    mean,
+    rstd,
+    eps,
+    D: tl.constexpr,
     BD: tl.constexpr,
     ACTIVATION: tl.constexpr,
     IS_RMS_NORM: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_BIAS: tl.constexpr,
 ):
+    i_t = tl.program_id(0).to(tl.int64)
     x += i_t * D
     y += i_t * D
     g += i_t * D
     elif ACTIVATION == "sigmoid":
         b_y = b_y * tl.sigmoid(b_g)
     tl.store(y + o_d, b_y, mask=m_d)
     key=["D", "NB", "IS_RMS_NORM", "HAS_DRESIDUAL", "HAS_WEIGHT"],
     **autotune_cache_kwargs,
 )
+@triton.jit(do_not_specialize=['T'])
 def layer_norm_gated_bwd_kernel(
+    x,
+    g,
+    w,
+    b,
+    y,
+    dy,
+    dx,
+    dg,
+    dw,
+    db,
     dresidual,
     dresidual_in,
     mean,
     HAS_BIAS: tl.constexpr,
     RECOMPUTE_OUTPUT: tl.constexpr,
 ):
+    i_s = tl.program_id(0).to(tl.int64)
     o_d = tl.arange(0, BD)
     m_d = o_d < D
     if HAS_WEIGHT:
         b_b = tl.load(b + o_d, mask=m_d, other=0.0).to(tl.float32)
         b_db = tl.zeros((BT, BD), dtype=tl.float32)
     for i_t in range(i_s * BS, i_s * BS + BS, BT):
+        o_t = (i_t + tl.arange(0, BT)).to(tl.int64)
+        m_t = o_t < T
+        m_x = m_t[:, None] & m_d[None, :]
+        p_x = x + o_t[:, None] * D + o_d[None, :]
+        p_g = g + o_t[:, None] * D + o_d[None, :]
+        p_dy = dy + o_t[:, None] * D + o_d[None, :]
+        p_dx = dx + o_t[:, None] * D + o_d[None, :]
+        p_dg = dg + o_t[:, None] * D + o_d[None, :]
         # [BT, BD]
+        b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32)
+        b_g = tl.load(p_g, mask=m_x, other=0.0).to(tl.float32)
+        b_dy = tl.load(p_dy, mask=m_x, other=0.0).to(tl.float32)
         if not IS_RMS_NORM:
+            p_mean = mean + o_t
+            b_mean = tl.load(p_mean, mask=m_t, other=0.0)
+        p_rstd = rstd + o_t
+        b_rstd = tl.load(p_rstd, mask=m_t, other=0.0)
         b_xhat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None]
         b_xhat = tl.where(m_d[None, :], b_xhat, 0.0)
         b_y = b_xhat * b_w[None, :] if HAS_WEIGHT else b_xhat
         if HAS_BIAS:
             b_y = b_y + b_b[None, :]
         b_sigmoid_g = tl.sigmoid(b_g)
         if ACTIVATION == "swish" or ACTIVATION == "silu":
+            b_gate = b_g * b_sigmoid_g
             b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
         elif ACTIVATION == "sigmoid":
+            b_gate = b_sigmoid_g
             b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
+        # b_dg needs pre-gate b_y; the recomputed output must include the gate to match the caller's linear input.
+        if RECOMPUTE_OUTPUT:
+            p_y = y + o_t[:, None] * D + o_d[None, :]
+            tl.store(p_y, (b_y * b_gate).to(p_y.dtype.element_ty), mask=m_x)
+        b_dy = b_dy * b_gate
         b_wdy = b_dy
         if HAS_WEIGHT or HAS_BIAS:
             b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D
             b_dx = (b_wdy - b_xhat * b_c1[:, None]) * b_rstd[:, None]
         if HAS_DRESIDUAL:
+            p_dres = dresidual + o_t[:, None] * D + o_d[None, :]
+            b_dres = tl.load(p_dres, mask=m_x, other=0.0).to(tl.float32)
             b_dx += b_dres
         if STORE_DRESIDUAL:
+            p_dres_in = dresidual_in + o_t[:, None] * D + o_d[None, :]
+            tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), mask=m_x)
+        tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=m_x)
+        tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_x)
     if HAS_WEIGHT:
         tl.store(dw + i_s * D + o_d, tl.sum(b_dw, axis=0), mask=m_d)
     key=["D", "IS_RMS_NORM", "STORE_DRESIDUAL", "HAS_DRESIDUAL", "HAS_WEIGHT"],
     **autotune_cache_kwargs,
 )
+@triton.jit(do_not_specialize=['T'])
+def layer_norm_gated_bwd_kernel_row(
+    x,
+    g,
+    w,
+    b,
+    y,
+    dy,
+    dx,
+    dg,
+    dw,
+    db,
     dresidual,
     dresidual_in,
     mean,
     HAS_BIAS: tl.constexpr,
     RECOMPUTE_OUTPUT: tl.constexpr,
 ):
+    i_s = tl.program_id(0).to(tl.int64)
     o_d = tl.arange(0, BD)
     mask = o_d < D
     x += i_s * BS * D
         b_db = tl.zeros((BD,), dtype=tl.float32)
     for i_t in range(i_s * BS, min(i_s * BS + BS, T)):
         b_x = tl.load(x + o_d, mask=mask, other=0).to(tl.float32)
         b_g = tl.load(g + o_d, mask=mask, other=0).to(tl.float32)
         b_dy = tl.load(dy + o_d, mask=mask, other=0).to(tl.float32)
         if not IS_RMS_NORM:
             b_mean = tl.load(mean + i_t)
         b_rstd = tl.load(rstd + i_t)
         b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
         b_xhat = tl.where(mask, b_xhat, 0.0)
         b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat
         if HAS_BIAS:
             b_y = b_y + b_b
         b_sigmoid_g = tl.sigmoid(b_g)
         if ACTIVATION == "swish" or ACTIVATION == "silu":
+            b_gate = b_g * b_sigmoid_g
             b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
         elif ACTIVATION == "sigmoid":
+            b_gate = b_sigmoid_g
             b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
+        # b_dg needs pre-gate b_y; the recomputed output must include the gate to match the caller's linear input.
+        if RECOMPUTE_OUTPUT:
+            tl.store(y + o_d, b_y * b_gate, mask=mask)
+        b_dy = b_dy * b_gate
         b_wdy = b_dy
         if HAS_WEIGHT:
             b_wdy = b_dy * b_w
         if HAS_DRESIDUAL:
             b_dres = tl.load(dresidual + o_d, mask=mask, other=0).to(tl.float32)
             b_dx += b_dres
         if STORE_DRESIDUAL:
             tl.store(dresidual_in + o_d, b_dx, mask=mask)
         tl.store(dx + o_d, b_dx, mask=mask)
         tl.store(db + i_s * D + o_d, b_db, mask=mask)
+# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support
+# @dispatch
 def layer_norm_gated_fwd(
     x: torch.Tensor,
     g: torch.Tensor,
     bias: torch.Tensor,
     activation: str = "swish",
     eps: float = 1e-5,
+    residual: torch.Tensor | None = None,
+    out_dtype: torch.dtype | None = None,
+    residual_dtype: torch.dtype | None = None,
     is_rms_norm: bool = False,
 ):
     if residual is not None:
         assert weight.shape == (D,)
     if bias is not None:
         assert bias.shape == (D,)
     y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
     if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype):
         residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
         residual_out = None
     mean = torch.empty((T,), dtype=torch.float, device=x.device) if not is_rms_norm else None
     rstd = torch.empty((T,), dtype=torch.float, device=x.device)
     MAX_FUSED_SIZE = 65536 // x.element_size()
     BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
     if D > BD:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
     if D <= 512:
+        # bucket token counts to limit autotuning across sequence lengths.
         NB = triton.cdiv(T, 2048 * 32)
         def grid(meta):
             IS_RMS_NORM=is_rms_norm,
         )
     else:
+        layer_norm_gated_fwd_kernel_row[(T,)](
             x=x,
             g=g,
             y=y,
     return y, mean, rstd, residual_out if residual_out is not None else x
+# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support
+# @dispatch
 def layer_norm_gated_bwd(
     dy: torch.Tensor,
     x: torch.Tensor,
     bias: torch.Tensor,
     activation: str = "swish",
     eps: float = 1e-5,
+    mean: torch.Tensor | None = None,
+    rstd: torch.Tensor | None = None,
+    dresidual: torch.Tensor | None = None,
     has_residual: bool = False,
     is_rms_norm: bool = False,
+    x_dtype: torch.dtype | None = None,
     recompute_output: bool = False,
 ):
     T, D = x.shape
         assert weight.shape == (D,)
     if bias is not None:
         assert bias.shape == (D,)
     dx = torch.empty_like(x) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
     dg = torch.empty_like(g) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device)
     dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None
     y = torch.empty(T, D, dtype=dy.dtype, device=dy.device) if recompute_output else None
     MAX_FUSED_SIZE = 65536 // x.element_size()
     BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
     if D > BD:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    # cap the program count for short sequences.
     NS = min(get_multiprocessor_count(x.device.index), T)
     BS = math.ceil(T / NS)
     grid = (NS,)
     if D <= 512:
+        # bucket token counts to limit autotuning across sequence lengths.
         NB = triton.cdiv(T, 2048 * 32)
         layer_norm_gated_bwd_kernel[grid](
             mean=mean,
             rstd=rstd,
             T=T,
             BS=BS,
+            D=D,
             BD=BD,
             NB=NB,
             ACTIVATION=activation,
             STORE_DRESIDUAL=dresidual_in is not None,
         )
     else:
+        layer_norm_gated_bwd_kernel_row[grid](
             x=x,
             g=g,
             w=weight,
             mean=mean,
             rstd=rstd,
             T=T,
             BS=BS,
+            D=D,
             BD=BD,
             ACTIVATION=activation,
             IS_RMS_NORM=is_rms_norm,
         )
     dw = dw.sum(0).to(weight.dtype) if weight is not None else None
     db = db.sum(0).to(bias.dtype) if bias is not None else None
+    # reuse dx when the residual gradient has the same dtype.
     if has_residual and dx.dtype == x.dtype:
         dresidual_in = dx
     return (dx, dg, dw, db, dresidual_in) if not recompute_output else (dx, dg, dw, db, dresidual_in, y)
     ):
         x_shape_og = x.shape
         g_shape_og = g.shape
         x = x.reshape(-1, x.shape[-1])
         g = g.reshape(-1, g.shape[-1])
         if residual is not None:
     ):
         x_shape_og = x.shape
         g_shape_og = g.shape
         x = x.reshape(-1, x.shape[-1])
         g = g.reshape(-1, g.shape[-1])
         if residual is not None:
         linear_weight = linear_weight.to(dtype)
         linear_bias = linear_bias.to(dtype) if linear_bias is not None else None
         out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias)
+        # recompute y in backward to save memory.
         ctx.save_for_backward(residual_out, g, norm_weight, norm_bias, linear_weight, mean, rstd)
         ctx.x_shape_og = x_shape_og
         ctx.g_shape_og = g_shape_og
     residual_in_fp32: bool = False,
     eps: float = 1e-6,
 ):
+    return LayerNormGatedFunction.apply(x, g, weight, bias, activation, residual, eps, prenorm, residual_in_fp32, False)
 def rms_norm_gated(
     residual_in_fp32: bool = False,
     eps: float = 1e-6,
 ):
+    return LayerNormGatedFunction.apply(x, g, weight, bias, activation, residual, eps, prenorm, residual_in_fp32, True)
 def layer_norm_swish_gate_linear(
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return layer_norm_gated(
+            x=x,
+            g=g,
+            weight=self.weight,
+            bias=self.bias,
+            activation=self.activation,
             residual=residual,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
+            eps=self.eps,
         )
         activation: str = "swish",
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return rms_norm_gated(
+            x=x,
+            g=g,
+            weight=self.weight,
+            bias=self.bias,
+            activation=self.activation,
             residual=residual,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
+            eps=self.eps,
         )
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
         super().__init__(
             hidden_size=hidden_size,
             elementwise_affine=elementwise_affine,
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__(hidden_size=hidden_size, elementwise_affine=elementwise_affine, eps=eps, device=device, dtype=dtype)
 class FusedLayerNormGatedLinear(nn.Module):
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return layer_norm_swish_gate_linear(
+            x=x,
+            g=g,
+            norm_weight=self.weight,
+            norm_bias=self.bias,
+            linear_weight=weight,
+            linear_bias=bias,
             residual=residual,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
+            eps=self.eps,
         )
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__(hidden_size=hidden_size, elementwise_affine=elementwise_affine, eps=eps, device=device, dtype=dtype)
 class FusedRMSNormGatedLinear(nn.Module):
     def __init__(
         self,
+        hidden_size: int,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
         residual: torch.Tensor | None = None,
         prenorm: bool = False,
         residual_in_fp32: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return rms_norm_swish_gate_linear(
+            x=x,
+            g=g,
+            norm_weight=self.weight,
+            norm_bias=self.bias,
+            linear_weight=weight,
+            linear_bias=bias,
             residual=residual,
             prenorm=prenorm,
             residual_in_fp32=residual_in_fp32,
+            eps=self.eps,
         )
         eps: float = 1e-5,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__(hidden_size=hidden_size, elementwise_affine=elementwise_affine, eps=eps, device=device, dtype=dtype)
+__all__ = [
+    'FusedLayerNormGated',
+    'FusedLayerNormGatedLinear',
+    'FusedLayerNormSwishGate',
+    'FusedLayerNormSwishGateLinear',
+    'FusedRMSNormGated',
+    'FusedRMSNormGatedLinear',
+    'FusedRMSNormSwishGate',
+    'FusedRMSNormSwishGateLinear',
+    'layer_norm_gated',
+    'layer_norm_gated_bwd',
+    'layer_norm_gated_fwd',
+    'layer_norm_swish_gate_linear',
+    'rms_norm_gated',
+    'rms_norm_swish_gate_linear',
+]